最小生成樹之 Prim 演算法

labuladong發表於2022-01-28

讀完本文,你不僅學會了演算法套路,還可以順便去 LeetCode 上拿下如下題目:

1135. 最低成本聯通所有城市(中等)

1584. 連線所有點的最小費用(中等)

-----------

本文是第 7 篇圖論演算法文章,先列舉一下我之前寫過的圖論演算法:

1、圖論演算法基礎

2、二分圖判定演算法

3、環檢測和拓撲排序演算法

4、Dijkstra 最短路徑演算法

5、Union Find 並查集演算法

6、Kruskal 最小生成樹演算法

像圖論演算法這種高階演算法雖然不算難,但是閱讀量普遍比較低,我本來是不想寫 Prim 演算法的,但考慮到演算法知識結構的完整性,我還是想把 Prim 演算法的坑填上,這樣所有經典的圖論演算法就基本完善了。

Prim 演算法和 Kruskal 演算法都是經典的最小生成樹演算法,閱讀本文之前,希望你讀過前文 Kruskal 最小生成樹演算法,瞭解最小生成樹的基本定義以及 Kruskal 演算法的基本原理,這樣就能很容易理解 Prim 演算法邏輯了。

對比 Kruskal 演算法

圖論的最小生成樹問題,就是讓你從圖中找若干邊形成一個邊的集合 mst,這些邊有以下特性:

1、這些邊組成的是一棵樹(樹和圖的區別在於不能包含環)。

2、這些邊形成的樹要包含所有節點。

3、這些邊的權重之和要儘可能小。

那麼 Kruskal 演算法是使用什麼邏輯滿足上述條件,計算最小生成樹的呢?

首先,Kruskal 演算法用到了貪心思想,來滿足權重之和儘可能小的問題:

先對所有邊按照權重從小到大排序,從權重最小的邊開始,選擇合適的邊加入 mst 集合,這樣挑出來的邊組成的樹就是權重和最小的。

其次,Kruskal 演算法用到了 Union-Find 並查集演算法,來保證挑選出來的這些邊組成的一定是一棵「樹」,而不會包含環或者形成一片「森林」:

如果一條邊的兩個節點已經是連通的,則這條邊會使樹中出現環;如果最後的連通分量總數大於 1,則說明形成的是「森林」而不是一棵「樹」。

那麼,本文的主角 Prim 演算法是使用什麼邏輯來計算最小生成樹的呢?

首先,Prim 演算法也使用貪心思想來讓生成樹的權重儘可能小,也就是「切分定理」,這個後文會詳細解釋。

其次,Prim 演算法使用 BFS 演算法思想visited 布林陣列避免成環,來保證選出來的邊最終形成的一定是一棵樹。

Prim 演算法不需要事先對所有邊排序,而是利用優先順序佇列動態實現排序的效果,所以我覺得 Prim 演算法類似於 Kruskal 的動態過程。

下面介紹一下 Prim 演算法的核心原理:切分定理。

切分定理

「切分」這個術語其實很好理解,就是將一幅圖分為兩個不重疊且非空的節點集合:

紅色的這一刀把圖中的節點分成了兩個集合,就是一種「切分」,其中被紅線切中的的邊(標記為藍色)叫做「橫切邊」。

PS:記住這兩個專業術語的意思,後面我們會頻繁使用這兩個詞,別搞混了。

當然,一幅圖肯定可以有若干種切分,因為根據切分的定義,只要你能一刀把節點分成兩部分就行。

接下來我們引入「切分定理」:

對於任意一種「切分」,其中權重最小的那條「橫切邊」一定是構成最小生成樹的一條邊

這應該很容易證明,如果一幅加權無向圖存在最小生成樹,假設下圖中用綠色標出來的邊就是最小生成樹:

那麼,你肯定可以找到若干「切分」方式,將這棵最小生成樹切成兩棵子樹。比如下面這種切分:

你會發現,任選一條藍色的「橫切邊」都可以將這兩棵子樹連線起來,構成一棵生成樹。

那麼為了讓最終這棵生成樹的權重和最小,你說你要怎麼選?

肯定選權重最小的那條「橫切邊」對吧,這就證明了切分定理。

關於切分定理,你也可以用反證法證明:

給定一幅圖的最小生成樹,那麼隨便給一種「切分」,一定至少有一條「橫切邊」屬於最小生成樹。

假設這條「橫切邊」不是權重最小的,那說明最小生成樹的權重和就還有再減小的餘地,那這就矛盾了,最小生成樹的權重和本來就是最小的,怎麼再減?所以切分定理是正確的。

有了這個切分定理,你大概就有了一個計算最小生成樹的演算法思路了:

既然每一次「切分」一定可以找到最小生成樹中的一條邊,那我就隨便切唄,每次都把權重最小的「橫切邊」拿出來加入最小生成樹,直到把構成最小生成樹的所有邊都切出來為止

嗯,可以說這就是 Prim 演算法的核心思路,不過具體實現起來,還是要有些技巧的。

因為你沒辦法讓計算機理解什麼叫「隨便切」,所以應該設計機械化的規則和章法來調教你的演算法,並儘量減少無用功。

Prim 演算法實現

我們思考演算法問題時,如果問題的一般情況不好解決,可以從比較簡單的特殊情況入手,Prim 演算法就是使用的這種思路。

按照「切分」的定義,只要把圖中的節點切成兩個不重疊且非空的節點集合即可算作一個合法的「切分」,那麼我只切出來一個節點,是不是也算是一個合法的「切分」?

是的,這是最簡單的「切分」,而且「橫切邊」也很好確定,就是這個節點的邊。

那我們就隨便選一個點,假設就從 A 點開始切分:

既然這是一個合法的「切分」,那麼按照切分定理,這些「橫切邊」AB, AF 中權重最小的邊一定是最小生成樹中的一條邊:

好,現在已經找到最小生成樹的第一條邊(邊 AB),然後呢,如何安排下一次「切分」?

按照 Prim 演算法的邏輯,我們接下來可以圍繞 AB 這兩個節點做切分:

然後又可以從這個切分產生的橫切邊(圖中藍色的邊)中找出權重最小的一條邊,也就又找到了最小生成樹中的第二條邊 BC

接下來呢?也是類似的,再圍繞著 A, B, C 這三個點做切分,產生的橫切邊中權重最小的邊是 BD,那麼 BD 就是最小生成樹的第三條邊:

接下來再圍繞 A, B, C, D 這四個點做切分……

Prim 演算法的邏輯就是這樣,每次切分都能找到最小生成樹的一條邊,然後又可以進行新一輪切分,直到找到最小生成樹的所有邊為止

這樣設計演算法有一個好處,就是比較容易確定每次新的「切分」所產生的「橫切邊」。

比如回顧剛才的圖,當我知道了節點 A, B 的所有「橫切邊」(不妨表示為 cut({A, B})),也就是圖中藍色的邊:

是否可以快速算出 cut({A, B, C}),也就是節點 A, B, C 的所有「橫切邊」有哪些?

是可以的,因為我們發現:

cut({A, B, C}) = cut({A, B}) + cut({C})

cut({C}) 就是節點 C 的所有鄰邊:

這個特點使我們用我們寫程式碼實現「切分」和處理「橫切邊」成為可能:

在進行切分的過程中,我們只要不斷把新節點的鄰邊加入橫切邊集合,就可以得到新的切分的所有橫切邊。

當然,細心的讀者肯定發現了,cut({A, B}) 的橫切邊和 cut({C}) 的橫切邊中 BC 邊重複了。

不過這很好處理,用一個布林陣列 inMST 輔助,防止重複計算橫切邊就行了。

最後一個問題,我們求橫切邊的目的是找權重最小的橫切邊,怎麼做到呢?

很簡單,用一個優先順序佇列儲存這些橫切邊,就可以動態計算權重最小的橫切邊了。

明白了上述演算法原理,下面來看一下 Prim 演算法的程式碼實現

class Prim {
    // 核心資料結構,儲存「橫切邊」的優先順序佇列
    private PriorityQueue<int[]> pq;
    // 類似 visited 陣列的作用,記錄哪些節點已經成為最小生成樹的一部分
    private boolean[] inMST;
    // 記錄最小生成樹的權重和
    private int weightSum = 0;
    // graph 是用鄰接表表示的一幅圖,
    // graph[s] 記錄節點 s 所有相鄰的邊,
    // 三元組 int[]{from, to, weight} 表示一條邊
    private List<int[]>[] graph;

    public Prim(List<int[]>[] graph) {
        this.graph = graph;
        this.pq = new PriorityQueue<>((a, b) -> {
            // 按照邊的權重從小到大排序
            return a[2] - b[2];
        });
        // 圖中有 n 個節點
        int n = graph.length;
        this.inMST = new boolean[n];

        // 隨便從一個點開始切分都可以,我們不妨從節點 0 開始
        inMST[0] = true;
        cut(0);
        // 不斷進行切分,向最小生成樹中新增邊
        while (!pq.isEmpty()) {
            int[] edge = pq.poll();
            int to = edge[1];
            int weight = edge[2];
            if (inMST[to]) {
                // 節點 to 已經在最小生成樹中,跳過
                // 否則這條邊會產生環
                continue;
            }
            // 將邊 edge 加入最小生成樹
            weightSum += weight;
            inMST[to] = true;
            // 節點 to 加入後,進行新一輪切分,會產生更多橫切邊
            cut(to);
        }
    }

    // 將 s 的橫切邊加入優先佇列
    private void cut(int s) {
        // 遍歷 s 的鄰邊
        for (int[] edge : graph[s]) {
            int to = edge[1];
            if (inMST[to]) {
                // 相鄰接點 to 已經在最小生成樹中,跳過
                // 否則這條邊會產生環
                continue;
            }
            // 加入橫切邊佇列
            pq.offer(edge);
        }
    }

    // 最小生成樹的權重和
    public int weightSum() {
        return weightSum;
    }

    // 判斷最小生成樹是否包含圖中的所有節點
    public boolean allConnected() {
        for (int i = 0; i < inMST.length; i++) {
            if (!inMST[i]) {
                return false;
            }
        }
        return true;
    }
}

明白了切分定理,加上詳細的程式碼註釋,你應該能夠看懂 Prim 演算法的程式碼了。

這裡我們可以再回顧一下本文開頭說的 Prim 演算法和 Kruskal 演算法 的聯絡:

Kruskal 演算法是在一開始的時候就把所有的邊排序,然後從權重最小的邊開始挑選屬於最小生成樹的邊,組建最小生成樹。

Prim 演算法是從一個起點的切分(一組橫切邊)開始執行類似 BFS 演算法的邏輯,藉助切分定理和優先順序佇列動態排序的特性,從這個起點「生長」出一棵最小生成樹。

說到這裡,Prim 演算法的時間複雜度是多少呢

這個不難分析,複雜度主要在優先順序佇列 pq 的操作上,由於 pq 裡面裝的是圖中的「邊」,假設一幅圖邊的條數為 E,那麼最多操作 O(E)pq。每次操作優先順序佇列的時間複雜度取決於佇列中的元素個數,取最壞情況就是 O(logE)

所以這種 Prim 演算法實現的總時間複雜度是 O(ElogE)。回想一下 Kruskal 演算法,它的時間複雜度主要是給所有邊按照權重排序,也是 O(ElogE)

不過話說回來,和前文 Dijkstra 演算法 類似,Prim 演算法的時間複雜度也是可以優化的,但優化點在於優先順序佇列的實現上,和 Prim 演算法本身的演算法思想關係不大,所以我們這裡就不做討論了,有興趣的讀者可以自行搜尋。

接下來,我們實操一波,把之前用 Kruskal 演算法解決的力扣題目運用 Prim 演算法再解決一遍。

題目實踐

第一題是力扣第 1135 題「最低成本聯通所有城市」,這是一道標準的最小生成樹問題:

函式簽名如下:

int minimumCost(int n, int[][] connections);

每座城市相當於圖中的節點,連通城市的成本相當於邊的權重,連通所有城市的最小成本即是最小生成樹的權重之和。

那麼解法就很明顯了,我們先把題目輸入的 connections 轉化成鄰接表形式,然後輸入給之前實現的 Prim 演算法類即可:

public int minimumCost(int n, int[][] connections) {
    // 轉化成無向圖鄰接表的形式
    List<int[]>[] graph = buildGraph(n, connections);
    // 執行 Prim 演算法
    Prim prim = new Prim(graph);

    if (!prim.allConnected()) {
        // 最小生成樹無法覆蓋所有節點
        return -1;
    }

    return prim.weightSum();
}

List<int[]>[] buildGraph(int n, int[][] connections) {
    // 圖中共有 n 個節點
    List<int[]>[] graph = new LinkedList[n];
    for (int i = 0; i < n; i++) {
        graph[i] = new LinkedList<>();
    }
    for (int[] conn : connections) {
        // 題目給的節點編號是從 1 開始的,
        // 但我們實現的 Prim 演算法需要從 0 開始編號
        int u = conn[0] - 1;
        int v = conn[1] - 1;
        int weight = conn[2];
        // 「無向圖」其實就是「雙向圖」
        // 一條邊表示為 int[]{from, to, weight}
        graph[u].add(new int[]{u, v, weight});
        graph[v].add(new int[]{v, u, weight});
    }
    return graph;
}

class Prim { /* 見上文 */ }

關於 buildGraph 函式需要注意兩點:

一是題目給的節點編號是從 1 開始的,所以我們做一下索引偏移,轉化成從 0 開始以便 Prim 類使用;

二是如何用鄰接表表示無向加權圖,前文 圖論演算法基礎 說過「無向圖」其實就可以理解為「雙向圖」。

這樣,我們轉化出來的 graph 形式就和之前的 Prim 演算法類對應了,可以直接施展 Prim 演算法計算最小生成樹。

再來看看力扣第 1584 題「連線所有點的最小費用」:

比如題目給的例子:

points = [[0,0],[2,2],[3,10],[5,2],[7,0]]

演算法應該返回 20,按如下方式連通各點:

函式簽名如下:

int minCostConnectPoints(int[][] points);

很顯然這也是一個標準的最小生成樹問題:每個點就是無向加權圖中的節點,邊的權重就是曼哈頓距離,連線所有點的最小費用就是最小生成樹的權重和。

所以我們只要把 points 陣列轉化成鄰接表的形式,即可複用之前實現的 Prim 演算法類:

public int minCostConnectPoints(int[][] points) {
    int n = points.length;
    List<int[]>[] graph = buildGraph(n, points);
    return new Prim(graph).weightSum();
}

// 構造無向圖
List<int[]>[] buildGraph(int n, int[][] points) {
    List<int[]>[] graph = new LinkedList[n];
    for (int i = 0; i < n; i++) {
        graph[i] = new LinkedList<>();
    }
    // 生成所有邊及權重
    for (int i = 0; i < n; i++) {
        for (int j = i + 1; j < n; j++) {
            int xi = points[i][0], yi = points[i][1];
            int xj = points[j][0], yj = points[j][1];
            int weight = Math.abs(xi - xj) + Math.abs(yi - yj);
            // 用 points 中的索引表示座標點
            graph[i].add(new int[]{i, j, weight});
            graph[j].add(new int[]{j, i, weight});
        }
    }
    return graph;
}

class Prim { /* 見上文 */ }

這道題做了一個小的變通:每個座標點是一個二元組,那麼按理說應該用五元組表示一條帶權重的邊,但這樣的話不便執行 Prim 演算法;所以我們用 points 陣列中的索引代表每個座標點,這樣就可以直接複用之前的 Prim 演算法邏輯了。

到這裡,Prim 演算法就講完了,整個圖論演算法也整的差不多了,更多精彩文章,敬請期待。

相關文章