Prim 最小生成树算法 :: labuladong的算法小抄

[utterances-bot]

文章链接点这里：Prim 最小生成树算法

评论礼仪 见这里，违者直接拉黑。

[feifeijin]

// 图中有 n 个节点
int n = graph.length;

这里的n不是节点数量，而是存在的边的数量。

[feifeijin]

我错了，n就是节点数量。。。可我不知道如何删除评论

[Alwin4Zhang]

1135 题「 最低成本联通所有城市」，Prim算法的python实现

import heapq
import collections
class Solution:
    # Kruscal并查集方式判定
    class UnionFind(object):
        def __init__(self, n):
            self.parent = list(range(n))

        def find(self, x):
            while self.parent[x] != x:
                self.parent[x] = self.parent[self.parent[x]]
                x = self.parent[x]
            return x

        def union(self, x1, x2):
            self.parent[self.find(x1)] = self.find(x2)

        def connected(self, x1, x2):
            return self.find(x1) == self.find(x2)

    def minCostConnectPoints(self, points: List[List[int]]) -> int:
        # 连边权重换成了曼哈顿距离
        def manhattan_distance(p1, p2):
            return abs(p1[0] - p2[0]) + abs(p1[1] - p2[1])

        # 1.uf并查集解决方案
        n = len(points)
        # uf = Solution.UnionFind(n)
        # # 遍历构建点之间距离数组
        distances = []
        for i in range(n):
            for j in range(i, n):
                distances.append(
                    (i, j, manhattan_distance(points[i], points[j])))

        # distances = sorted(distances, key=lambda x: x[-1])

        # count = 0
        # for dist in distances:
        #     x, y, x_y_dist = dist
        #     if uf.connected(x, y):
        #         continue
        #     uf.union(x, y)
        #     count += x_y_dist
        # return count

        # 2.BFS方式
        def build_graph(connections):  # 构建图
            graph = collections.defaultdict(list)
            # 无向图就是双向图
            for conn in connections:
                u, v, weight = conn
                graph[u].append((u, v, weight))
                graph[v].append((v, u, weight))
            return graph

        def cut(s):
            for edge in graph[s]:
                _from, _to, weight = edge
                if in_mst[_to]:
                    continue
                heapq.heappush(pq, (weight, _from, _to))

        graph = build_graph(distances)
        pq = []
        # 记录MST中的节点
        in_mst = [False] * n
        # 记录最终结果
        weight_sum = 0
        in_mst[0] = True
        cut(0)

        while pq:
            weight, _from, _to = heapq.heappop(pq)
            if in_mst[_to]:
                continue
            weight_sum += weight
            in_mst[_to] = True
            cut(_to)
        return weight_sum

[yonggui-yan]

棒。

注释 “# Kruscal并查集方式判定” 让我恍惚了。

[ZackHongru]

东哥 什么时候能有c++版的答案呢？
因为没有太多代码基础，看你的代码在写出c++版的有些吃力...

[deweicai23]

什么时候出个弗洛依德算法呢！！

[labuladong]

@ZackHongru 算法用的标准库就那么几个，看多了就习惯了。另外可以在评论区看看有没有别人的 cpp 代码。

[JerryWuZiJie]

“Prim 算法的时间复杂度也是可以优化的，但优化点在于优先级队列的实现上” 能优于O(ElogE)吗？网上好像都是从array优化成priorityqueue，好像没有更优化的了

[labuladong]

@JerryWuZiJie priorityqueue 也有不同的实现，有些实现可以直接修改队列中的元素，这样的优先级队列可以进一步降低时间复杂度。Java 标准库的优先级队列不可以。

[825344491]

对Prim算法的个人理解

个人认为，对于Prim算法的理解，用教科书上的“加点”可以更好接受。设已被连入最小生成树的集合位MST，则每一次选取连接到MST中cost最小的点。因此优先队列里永远放的都是点和该点连接到MST需要的最小cost（其实也还是边的权值），然后随着算法的执行，每个点连接到MST的cost也会不断被更新（不断往优先级队列里插入更小的cost，原来的cost不会被删除）。这样想可以把Prim和Dijkstra完全类比，两个算法只有一行代码不同。

附1135和1584两题Python简洁Prim解法：

1135. 最低成本联通所有城市

class Solution:
    def minimumCost(self, n: int, connections: List[List[int]]) -> int:
        graph = [[] for _ in range(n)]
        for connection in connections:
            graph[connection[0] - 1].append((connection[2], connection[1] - 1))
            graph[connection[1] - 1].append((connection[2], connection[0] - 1))
        heap = [(0, 0)]
        mst = set()
        cost = 0
        while heap:
            cur = heapq.heappop(heap)
            if cur[1] not in mst:
                mst.add(cur[1])
                cost += cur[0]
                for vertex in graph[cur[1]]:
                    heapq.heappush(heap, vertex)
        if len(mst) == n:
            return cost
        return -1

1584. 连接所有点的最小费用

class Solution:
    def minCostConnectPoints(self, points: List[List[int]]) -> int:
        n = len(points)
        graph = [[] for _ in range(n)]
        for i in range(n - 1):
            for j in range(i + 1, n):
                graph[i].append((abs(points[i][0] - points[j][0]) + abs(points[i][1] - points[j][1]), j))
                graph[j].append((abs(points[i][0] - points[j][0]) + abs(points[i][1] - points[j][1]), i))
        heap = [(0, 0)]
        mst = set()
        cost = 0
        while heap:
            cur = heapq.heappop(heap)
            if cur[1] not in mst:
                mst.add(cur[1])
                cost += cur[0]
                for vertex in graph[cur[1]]:
                    heapq.heappush(heap, vertex)
            if len(mst) == n:
                return cost

[yonggui-yan]

好简洁！厉害了！

[labuladong]

@825344491 👍

[amark1235]

cut({A, B, C}) = cut({A, B}) + cut({C}) - [cut({A,B}) ^cut({C}) ]，这样才对吧，减去它俩的交集, 刚写错了

[labuladong]

伪码，理解到意思就行

[WuuG]

cpp

1135.最低成本联通所有城市

    struct cmp
    {
        bool operator()(const vector<int> &edge1, const vector<int> &edge2)
        {
            return edge1[2] > edge2[2];
        }
    };

    class Solution
    {
        vector<bool> visited;
        vector<vector<vector<int>>> graph;
        using PQ = priority_queue<vector<int>, vector<vector<int>>, cmp>;
        PQ q;

    public:
        int minimumCost(int n, vector<vector<int>> &connections)
        {
            // 编号从 1 ~ n 为了直接使用序号，多了一个0节点，需要手动设置为true
            visited = vector<bool>(n + 1);
            visited[0] = true;
            buildGraph(n, connections);
            addPointsAndEdges(1);
            int sum = 0;
            while (!q.empty())
            {
                auto curEdge = q.top();
                int from = curEdge[0], to = curEdge[1], cost = curEdge[2];
                q.pop();
                // 访问过，说明该点已经连接
                if (visited[to])
                {
                    continue;
                }
                sum += cost;
                addPointsAndEdges(to);
            }
            for_each(visited.begin(), visited.end(), [&sum](bool isVisited)
                     {
            if(!isVisited)
            {
                sum = -1;
            } });
            return sum;
        }
        void buildGraph(int n, vector<vector<int>> &connections)
        {
            graph = vector<vector<vector<int>>>(n + 1);
            for (auto connect : connections)
            {
                int from = connect[0], to = connect[1], cost = connect[2];
                graph[from].emplace_back(vector<int>{from, to, cost});
                graph[to].emplace_back(vector<int>{to, from, cost});
            }
        }
        void addPointsAndEdges(int cur)
        {
            visited[cur] = true;
            for (auto nxtEdge : graph[cur])
            {
                int from = nxtEdge[0], to = nxtEdge[1], cost = nxtEdge[2];
                if (visited[to])
                {
                    continue;
                }
                q.push(nxtEdge);
            }
        }
    };

1584.连接所有点的最小费用

    class Solution
    {
        using PQ = priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>>;
        // 表示点的数量
        int n;
        unordered_set<int> inMst;
        // { 点距离， 去向的点}
        PQ q;

    public:
        int minCostConnectPoints(vector<vector<int>> &points)
        {
            n = points.size();
            addPointAndEdges(0, points);

            int sum = 0;
            while (!q.empty())
            {
                pair<int, int> curEdge = q.top();
                q.pop();
                int cost = curEdge.first, to = curEdge.second;
                if (inMst.count(to))
                {
                    continue;
                }

                sum += cost;
                addPointAndEdges(to, points);
                if (inMst.size() == n)
                {
                    break;
                }
            }
            return sum;
        }
        void addPointAndEdges(int cur, vector<vector<int>> &points)
        {
            inMst.insert(cur);
            for (int i = 0; i < n; ++i)
            {
                // 该点已经加入了mst中，就不需要加入了
                if (inMst.count(i))
                {
                    continue;
                }
                int x1 = points[cur][0], y1 = points[cur][1],
                    x2 = points[i][0], y2 = points[i][1];
                int diff = abs(x1 - x2) + abs(y1 - y2);
                q.push({diff, i});
            }
        }
    };

[NoOneHero]

1584

func minCostConnectPoints(points [][]int) int {
	n := len(points)
	// 生成边权重
	edges := make([][][]int, n)
	for i := 0; i < len(points); i++ {
		for j := i + 1; j < len(points); j++ {
			value := abs(points[i][0]-points[j][0]) + abs(points[i][1]-points[j][1])
			edges[i] = append(edges[i], []int{i, j, value})
			edges[j] = append(edges[j], []int{j, i, value})
		}
	}

	var res int

	visited := make([]bool, n)
	var q PQ
	visited[0] = true
	for _, edge := range edges[0] {
		heap.Push(&q, edge)
	}
	for len(q) > 0 {
		edge := heap.Pop(&q).([]int)
		// 连接点
		to := edge[1]
		if visited[to] {
			continue
		}
		// 权重结算
		res += edge[2]
		// 标记已访问
		visited[to] = true
		for _, edge := range edges[to] {
			if !visited[edge[1]] {
				heap.Push(&q, edge)
			}
		}
	}
	return res
}

func abs(a int) int {
	if a < 0 {
		return -a
	}
	return a
}

type PQ [][]int

func (pq PQ) Len() int {
	return len(pq)
}

func (pq PQ) Swap(a, b int) {
	pq[a], pq[b] = pq[b], pq[a]
}

func (pq PQ) Less(a, b int) bool {
	return pq[a][2] < pq[b][2]
}

func (pq *PQ) Push(nodeInterface interface{}) {
	node := nodeInterface.([]int)
	*pq = append(*pq, node)
}

func (pq *PQ) Pop() interface{} {
	old := *pq
	lastNode := old[len(*pq)-1]
	*pq = old[:len(*pq)-1]
	return lastNode
}

[bruceeewong]

打卡！理论上Kruskal和Prim时间复杂度一致，为啥leetcode上Prim要比Kruskal慢呢

[labuladong]

理论复杂度相同并不代表实际运行时间也相同。实际运行时间还会受到测试数据，代码的复杂程度等等多方面因素的影响。

[bugcat9]

1548题C++中使用prim算法超时了，不知道为啥。希望有大佬帮忙看看

struct cmp
{
    bool operator()(const vector<int> &edge1, const vector<int> &edge2)
    {
        // 按照边的权重从小到大排序
        return edge1[2] > edge2[2];
    }
};

class Prim
{
private:
    // 核心数据结构，存储「横切边」的优先级队列
    using PQ = priority_queue<vector<int>, vector<vector<int>>, cmp>;
    PQ pq;
    // 类似 visited 数组的作用，记录哪些节点已经成为最小生成树的一部分
    vector<bool> inMST;
    // 记录最小生成树的权重和
    int weightSum;
    // graph 是用邻接表表示的一幅图，
    // graph[s] 记录节点 s 所有相邻的边，
    // 三元组 int[]{from, to, weight} 表示一条边
    vector<vector<vector<int>>> graph;

    /**
     * @brief 将 s 的横切边加入优先队列
     *
     * @param s
     */
    void cut(int s)
    {
        for (auto &edge : graph[s])
        {
            int to = edge[1];
            if (inMST[to])
            {
                continue;
            }
            pq.push(edge);
        }
    }

public:
    Prim(const vector<vector<vector<int>>> &graph) : graph(graph)
    {
        int n = graph.size();
        inMST.resize(n);
        weightSum = 0;
        // 随便从一个点开始切分都可以，我们不妨从节点 0 开始
        inMST[0] = true;
        cut(0);
        // 不断进行切分，向最小生成树中添加边
        while (!pq.empty())
        {
            vector<int> edge = pq.top();
            pq.pop();
            int to = edge[1];
            int weight = edge[2];
            if (inMST[to])
            {
                // 节点 to 已经在最小生成树中，跳过
                // 否则这条边会产生环
                continue;
            }
            // 将边 edge 加入最小生成树
            weightSum += weight;
            inMST[to] = true;
            // 节点 to 加入后，进行新一轮切分，会产生更多横切边
            cut(to);
        }
    }

    /**
     * @brief 最小生成树的权重和
     * 
     * @return int 
     */
    int getWeightSum()
    {
        return weightSum;
    }

    /**
     * @brief 判断最小生成树是否包含图中的所有节点
     * 
     * @return true 
     * @return false 
     */
    bool allConnected()
    {
        for (int i = 0; i < inMST.size(); i++)
        {
            if (!inMST[i])
            {
                return false;
            }
        }
        return true;
    }
};

1584中代码

class Solution
{
public:
	int minCostConnectPoints(vector<vector<int>> &points)
	{
		int n = points.size();
		vector<vector<vector<int>>> graph = buildGraph(n, points);
		return Prim(graph).getWeightSum();
	}

	vector<vector<vector<int>>> buildGraph(int n, vector<vector<int>> &points)
	{
		vector<vector<vector<int>>> graph(n);
		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];
				// 用坐标点在 points 中的索引表示坐标点
				graph[i].push_back({i, j, abs(xi - xj) + abs(yi - yj)});
				graph[j].push_back({j, i, abs(xi - xj) + abs(yi - yj)});
			}
		}
		return graph;
	}
};

[wozhenbang2004]

我也是,照着他的kruskal 写也是超时的.你现在有思路了吗

[shishuochen]

个人感觉是队列操作的代码导致的。
我看楼主的process函数中再加入新的横切边时，会有一些剪枝

if (!visited[v] && w < dist[v]) {
                    dist[v] = w;
                    pq.push({dist[v], v});
}

我没仔细去读完整代码，但我觉得这个剪枝操作比较重要，因为我自己写的剪枝中，只判断了节点是否在树中，自然也TLE了，所以我改写了下while的终止条件，用树中节点数目来代替，过了，但是效率很低。Po个代码吧，写的有点乱，但是整体思路都在。

class Solution {
public:
    int minCostConnectPoints(vector<vector<int>>& points) {
        // 节点数目
        int pnums=points.size();
        // 计算距离矩阵: 无向图邻接表
        vector<vector<pair<int,int>>> matrix(pnums);
        for(int i=0;i<pnums-1;i++) {
            for(int j=i+1;j<pnums;j++) {
                // 两点之间的距离
                int x1= points[i][0]-points[j][0];
                int y1=points[i][1]-points[j][1];
                int xlim=x1>0 ? x1 : -x1;
                int ylim= y1>0 ? y1 : -y1;
                matrix[i].push_back(make_pair(j,xlim+ylim));
                matrix[j].push_back(make_pair(i,xlim+ylim));
            }
        }
        // 构建邻接表
        // 数据类型，容器，排序算法
        // {point i, point j, weight}
        auto cmp = [](const std::vector<int>& a, const std::vector<int>& b) { return a[2] > b[2]; };
    std::priority_queue<std::vector<int>, std::vector<std::vector<int>>, decltype(cmp)> Q(cmp);
        // 访问数组
        vector<bool> visited(pnums,false);
        visited[0]=true;
        // 从零开始
        for(auto s:matrix[0]) {
            Q.push({0,s.first,s.second});
        }
        int weight=0;
        int count=1;   // 树中的节点数目
        while(count<pnums) {
            vector<int> tt=Q.top();
            Q.pop();
            // 如果to已经在树中，则放弃
            int n1=tt[0],n2=tt[1],w=tt[2];
            if(visited[n2]) continue;
            visited[n2]=true;
            count++;
            weight+=w;   // 增加weight
            // 加入新的横切边
            for(auto s:matrix[n2]) {
                if(visited[s.first])
                    continue;
                Q.push({n2,s.first,s.second});
            }
        }
        return weight;
    }
};

[max870701]

Python中可用 heapq 作為 Priority Queue
1584. Min Cost to Connect All Points

import heapq
class Solution(object):
    def minCostConnectPoints(self, points):
        """
        :type points: List[List[int]]
        :rtype: int
        """
        n = len(points)
        graph = self.buildGraph(n, points)
        return Prim(graph).weightSum
    
    def buildGraph(self, n, points):
        graph = [[] for _ in range(n)]
        # 生成所有邊及權重  
        for i in range(n):
            for j in range(i+1, n):
                xi, yi = points[i][0], points[i][1]
                xj, yj = points[j][0], points[j][1]
                weight = abs(xi - xj) + abs(yi - yj)
                # 用 points 中的索引表示座標點
                graph[i].append((i, j, weight))
                graph[j].append((j, i, weight))
        
        return graph

class Prim(object):
    def __init__(self, graph):
        self.graph = graph
        n = len(graph)
        self.pq = []
        # Visited 數組功用
        self.inMST = [False for _ in range(n)]
        self.weightSum = 0
        # 隨便從一個節點開始分割，不妨從節點 0 開始
        self.inMST[0] = True
        self.cut(0)
        # 不斷進行分割，向最小生成樹中添加邊
        while self.pq:
            edge = heapq.heappop(self.pq)
            # 形式 (priority_order, (node1, node2, weight))
            to = edge[1][1]
            weight = edge[1][2]
            # 節點 to 已經在最小生成樹中，則跳過
            # 避免產生 Cycle
            if self.inMST[to]:
                continue
            # 將 edge 加入最小生成樹
            self.weightSum += weight
            self.inMST[to] = True
            self.cut(to)
    
    def cut(self, s):
        # 遍歷 s 的鄰邊
        for edge in self.graph[s]:
            # edge的終點
            to = edge[1]
            # 若相鄰接點 to 已經在最小生成樹中，跳過
            if self.inMST[to]:
                continue
            # 不在最小生成樹中，則加入優先級隊列(照權重)
            # heapq.heappush(heap, (priority order, value))
            heapq.heappush(self.pq, (edge[2], edge))

    def allConnected(self):
        # 判斷最小生成樹是否包含圖中所有節點
        for i in range(len(self.inMST)):
            if not self.inMST[i]:
                return False
        return True

[max870701]

錯了...是 1584. Min Cost to Connect All Points

[Yongfanchen]

c++的代码好像有问题。