2022-06-07 / Bartłomiej Kurek
Algorytm Prima: minimalne drzewo rozpinające

W tym artykule zajmiemy się implementacją zachłannego algorytmu Prima, którego zadaniem jest wyznaczenie minimalnego drzewa rozpinającego (minimum spanning tree).
Omawiany kod znajdziemy w repozytorium git.

Wikipedia: Prim's algorithm

Założenia implementacji

Algorytm operuje na nieskierowanym grafie ważonym. Każdej krawędzi przypiszemy zatem koszt, a same krawędzie zapiszemy w tabli edges. W algorytmie wykorzystywana jest kolejka priorytetowa, którą zrealizujemy jako osobną tabelę w bazie danych: prim_priority_queue. Wynikiem algorytmu (funkcji plpgsql) będzie minimalny koszt drzewa rozpinającego, a poszczególne krawędzie algorytm umieści w tabeli prim_spanning_tree.

Schemat bazy danych

Wierzchołki grafu będziemy przechowywać w tabeli edges:

CREATE TABLE edges (
       edge_id SERIAL PRIMARY KEY,
       src INTEGER NOT NULL,
       dst INTEGER NOT NULL,
       cost INTEGER NOT NULL
);

Tabela realizująca zadania kolejki priorytetowej: prim_priority_queue.

CREATE TABLE prim_priority_queue(
       pq_id SERIAL PRIMARY KEY,
       src INTEGER NOT NULL,
       cost INTEGER NOT NULL
);

Tabela, w której znajdą się poszczególne krawędzie drzewa rozpinającego: prim_spanning_tree.

CREATE TABLE prim_spanning_tree(
       id SERIAL PRIMARY KEY,
       u INTEGER NOT NULL,
       v INTEGER NOT NULL CHECK(v != u),
       cost INTEGER NOT NULL
);

Poszczególne krawędzie mogą zostać wyznaczone tylko jeden raz, zatem do tabeli prim_spanning_tree dodajemy unikalny index:

CREATE UNIQUE INDEX prim_spanning_tree_uidx ON prim_spanning_tree(u, v);

Algorytm

Krótki opis implementacji

Algorytm rozpoczyna się od wskazanego wierzchołka, który od razu umieszczamy w kolejce priorytetowej.
Następnie z kolejki wybieramy (i usuwamy) kolejną krawędź według najniższego kosztu. Znalezioną krawędź umieszczamy w tabeli wyników, aktualizujemy koszt całkowity drzewa, a wierzchołki osiągalne z kolejnego wierzchołka dodajemy do kolejki priorytetowej. Niuansem implementacyjnym jest konieczność śledzenia wierzchołków, ktore zostały już przetworzone.

Funkcja PLPGSQL

Tworzymy funkcję "prim(integer)". Argumentem funkcji jest wierzchołek startowy.

CREATE FUNCTION prim(src_ INTEGER DEFAULT 0)
RETURNS INTEGER
AS $$
DECLARE
    min_span_tree_cost INTEGER = 0;

    dst_ INTEGER;
    cost_ INTEGER;

    next_ INTEGER; -- next node (from priority queue)

    pq_cost_ INTEGER;
    pq_id_ INTEGER;
    pq_len_ INTEGER;

    is_added_ BOOLEAN;
    added_ INTEGER[] = ARRAY[]::INTEGER[];
BEGIN
    TRUNCATE prim_priority_queue;
    TRUNCATE prim_spanning_tree;

    -- Start vertex.
    INSERT INTO prim_priority_queue(src, cost) VALUES(src_, 0);

    WHILE TRUE
    LOOP
        -- Check priority_queue length. Exit if empty.
        SELECT count(*) FROM prim_priority_queue INTO pq_len_;
        IF pq_len_ = 0 THEN
           EXIT;
        END IF;

        -- Fetch next element from priority queue.
        SELECT pq_id, src, cost FROM prim_priority_queue ORDER BY cost ASC LIMIT 1
               INTO pq_id_, next_, pq_cost_;

        -- Remove element from priority queue.
        DELETE FROM prim_priority_queue WHERE pq_id = pq_id_; 

        -- Was node already added?
        is_added_ := (next_ = ANY(added_)); 

        IF NOT is_added_ THEN
            -- Node was not added. Update the cost and mark node added.
            min_span_tree_cost := min_span_tree_cost + pq_cost_;
            added_ := array_append(added_, next_);

            IF src_ != next_ THEN
                INSERT INTO prim_spanning_tree(u, v, cost) VALUES(src_, next_, pq_cost_);
            END IF;

            FOR dst_, cost_ IN
                SELECT dst, cost FROM edges WHERE src = next_
            LOOP
                is_added_ := (dst_ = ANY(added_));
                IF NOT is_added_ THEN
                   INSERT INTO prim_priority_queue(src, cost) VALUES(dst_, cost_);
                END IF;
            END LOOP;

            src_ := next_;
        END IF;
    END LOOP;

    RETURN min_span_tree_cost;
END
$$
LANGUAGE plpgsql;

Uruchomienie algorytmu

Do sprawdzenia algorytmu przygotowujemy prosty graf.
(Pliki graph.txt oraz Makefile).

Graf jest nieskierowany, zatem w tabeli edges umieszczamy rekordy "a -- b" oraz "b -- a".

INSERT INTO edges(src, dst, cost) VALUES
    (0, 1, 4), (0, 2, 1), (0, 3, 5),
    (1, 0, 4), (1, 3, 2), (1, 4, 3), (1, 5, 3),
    (2, 0, 1), (2, 3, 2), (2, 4, 8),
    (3, 0, 5), (3, 1, 2), (3, 2, 2), (3, 4, 1),
    (4, 1, 3), (4, 2, 8), (4, 3, 1), (4, 5, 3),
    (5, 1, 3), (5, 4, 3);

Uzyskujemy 20 rekordów w tabeli edges:

test=> SELECT * FROM edges
 edge_id | src | dst | cost 
---------+-----+-----+------
       1 |   0 |   1 |    4
       2 |   0 |   2 |    1
       3 |   0 |   3 |    5
       4 |   1 |   0 |    4
       5 |   1 |   3 |    2
       6 |   1 |   4 |    3
       7 |   1 |   5 |    3
       8 |   2 |   0 |    1
       9 |   2 |   3 |    2
      10 |   2 |   4 |    8
      11 |   3 |   0 |    5
      12 |   3 |   1 |    2
      13 |   3 |   2 |    2
      14 |   3 |   4 |    1
      15 |   4 |   1 |    3
      16 |   4 |   2 |    8
      17 |   4 |   3 |    1
      18 |   4 |   5 |    3
      19 |   5 |   1 |    3
      20 |   5 |   4 |    3
(20 rows)

Uruchamiamy naszą funkcję, dla przykładu wierzchołkiem startowym niech będzie wierzchołek "2":

test=> SELECT * FROM prim(2);
 prim 
------
    9
(1 row)

Całkowity koszt minimalnego drzewa rozpinającego wynosi 9. Wierzchołki (oraz poszczególne koszty) możemy pobrać z tabeli prim_spanning_tree:

test=> SELECT * FROM prim_spanning_tree;
 id | u | v | cost 
----+---+---+------
  1 | 2 | 0 |    1
  2 | 0 | 3 |    2
  3 | 3 | 4 |    1
  4 | 4 | 1 |    2
  5 | 1 | 5 |    3
(5 rows)

Suma kosztów poszczególnych krawędzi faktycznie wynosi 9:

test=> SELECT SUM(cost) FROM prim_spanning_tree;
 sum 
-----
   9
(1 row)