Leetcode 最小生成树系列（2）

水资源分配优化

村里面一共有 $n$ 栋房子。我们希望通过建造水井和铺设管道来为所有房子供水。

对于每个房子 $i$，我们有两种可选的供水方案：一种是直接在房子内建造水井，成本为 $wells[i-1]$ （注意 $-1$，因为 索引从$0$开始 ）；另一种是从另一口井铺设管道引水，数组 $pipes$ 给出了在房子间铺设管道的成本，其中每个 $pipes[j]=[house{1}_j,house{2}_j,cos{t}_j]$ 代表用管道将 $house{1}_j$ 和 $house{2}_j$连接在一起的成本。连接是双向的。

请返回为所有房子都供水的最低总成本 。

示例：

输入： n = 3, wells = [1,2,2], pipes = [[1,2,1],[2,3,1]]
输出： 3
解释：
上图展示了铺设管道连接房屋的成本。最好的策略是在第一个房子里建造水井（成本为 1），然后将其他房子铺设管道连起来（成本为 2），所以总成本为 3。

提示：

$2<=n<=104$
$wells.length==n$
$0<=wells[i]<=105$
$1<=pipes.length<=104$
$pipes[j].length==3$
$1<=house1j,house2j<=n$
$0<=costj<=105$
$house{1}_j!=house{2}_j$

题解：

题目的核心技巧是引入虚拟节点，把“打井”也看作是建一条特殊的边：增加一个虚拟节点 0，表示“水源”。对每个房子 $i$，加一条边 $(0,i,wells[i-1])$，代表直接在该房子打井。原有 $pipes$ 数据就是房子之间的无向边。问题转化为：在包含 $0,1,2,\dots ,n$ 节点的图上，求一棵最小生成树（MST）。

MST 会自动帮我们决定：哪些房子直接打井（连接到 0 节点的边被选中）。哪些房子通过管道连接。

算法求解步骤

• 把所有虚拟边 (0, i, wells[i-1]) 加入。
• 把所有管道边 (house1, house2, cost) 加入。
• 排序
• 将所有边按 cost 从小到大排序。
• 并查集合并
• 初始化并查集，节点范围 [0, n]。
• 从最小成本边开始遍历，若两个端点不在同一集合，就合并并加到总成本中。
• 当选中的边数量达到 n（因为有 n+1 个节点）时结束。
• 输出结果
• 累加的边权和即为最小成本。

算法代码实现

class Solution {

    static class DSU {
        int[] parent;
        int[] size;

        DSU(int n) {
            parent = new int[n];
            size = new int[n];
            for (int i = 0; i < n; i++) {
                parent[i] = i;
                size[i] = 1;
            }
        }

        int find(int x) {
            while (x != parent[x]) {
                parent[x] = parent[parent[x]];
                x = parent[x];
            }
            return x;
        }

        boolean union(int a, int b) {
            int ra = find(a), rb = find(b);
            if (ra == rb) return false;
            if (size[ra] < size[rb]) {
                int t = ra; ra = rb; rb = t;
            }
            parent[rb] = ra;
            size[ra] += size[rb];
            return true;
        }
    }

    static class Edge {
        int u, v, w;
        Edge(int u, int v, int w) { this.u = u; this.v = v; this.w = w; }
    }

    /**
     * @param n 房子数量（1..n）
     * @param wells wells[i] 为在房子 i+1 打井的成本
     * @param pipes pipes[j] = [house1, house2, cost] 为铺设管道成本（双向）
     * @return 供水的最小总成本
     */
    public static int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
        List<Edge> edges = new ArrayList<>(n + (pipes == null ? 0 : pipes.length));

        // 虚拟边：(0, i, wells[i-1])，表示在房子 i 打井
        for (int i = 1; i <= n; i++) {
            edges.add(new Edge(0, i, wells[i - 1]));
        }
        // 管道边
        if (pipes != null) {
            for (int[] p : pipes) {
                edges.add(new Edge(p[0], p[1], p[2]));
            }
        }

        // 按成本升序排序
        edges.sort(Comparator.comparingInt(e -> e.w));

        DSU dsu = new DSU(n + 1); // 节点 0..n
        long total = 0L;
        int used = 0; // 选边计数，目标为 n（n+1 个节点的生成树）

        for (Edge e : edges) {
            if (dsu.union(e.u, e.v)) {
                total += e.w;
                used++;
                if (used == n) break;
            }
        }
        return (int) total;
    }
}

运行效果