LeetCode 1168. 水资源分配优化

2021-03-31 HGNULB 约 13 分钟

文章信息 6 个分类 7 个标签

分类

题目描述

✅ 1168. 水资源分配优化

题意分析

村里有 n 座房子（编号 1 到 n），要让每座房子都有水。有两种花钱的方式：在第 i 座房子处打一口井，花费 wells[i-1]；或者在两座房子之间铺一条管道，花费由 pipes 给出，管道可以把水从一座已有水的房子引到另一座。求让所有房子都有水的最小总花费。

「有水」的传递性是理解本题的关键：一座房子有水，可能因为它自己有井，也可能因为它通过管道（可以经过若干中转）连到了某座有井的房子。所以最终的付费方案里，每一个「连通块」内必须至少有一口井，块内其余房子靠管道相连。

这就暴露出难点：打井和铺管道是两种性质完全不同的花费——前者作用在单个点上，后者作用在一条边上。如果分开考虑，就要枚举「哪些房子打井」再对每个方案求连通代价，组合爆炸。

再看约束：n 上限 $10^4$，pipes 长度上限 $10^4$。这个规模允许 $O(m \log m)$ 的排序型算法，但不允许对打井方案做任何形式的枚举。同时「所有点连通、总代价最小」这几个字，配合「图 + 边权」的输入形式，已经把最小生成树的味道摆得很明显了——缺的只是把打井也变成一条边。

边界：pipes 可能为空，此时每座房子只能各自打井，答案是所有 wells 之和；同一对房子之间可能给出多条管道（重复边），也可能存在自环式的冗余输入，算法必须能容忍；pipes 给出的图不保证连通。

解法：最小生成树 + 虚拟节点

核心思路

先看暴力：枚举打井房屋的子集，剩下的房子必须通过管道连到某个有井的房子上，再对每种子集求最小连通代价。子集数 $2^n$，$n = 10^4$ 时完全不可行。瓶颈在于把「打井」和「铺管」当成了两类互不相干的决策，于是必须先定一类再算另一类。

关键观察是把两者统一：打井，本质上是「从水源引一条水管到这座房子」。既然如此，就凭空造一个编号为 0 的虚拟节点代表「水源 / 地下水」，把「在房子 i 打井、花费 wells[i-1]」重写成一条边：
\[(0,\ i,\ wells[i-1])\]
这一步之后，所有花费都变成了边权，图上有 n + 1 个点（0 号加上 n 座房子）、n + m 条边（n 条虚拟井边加上 m 条管道边）。

现在原问题等价于什么？「每座房子都有水」等价于「每座房子都能沿着已付费的边连到 0 号节点」，也就是新图上所有点连通；「总花费最小」就是边权和最小。这正是最小生成树的定义。

这个等价要双向说清楚。任何合法的供水方案，对应新图上一个连通所有点的边子集（每个原连通块里的那口井对应一条连到 0 的边）；反过来，新图上任何连通所有点的边子集，把连到 0 的边解读为打井、其余解读为铺管，就是一个合法方案且花费相同。既然二者一一对应且代价相等，最小生成树的权和就是答案。

剩下的就是求 MST。选 Kruskal：把显式边集按权升序排序，依次用并查集合并两端。合并成功说明它是当前连接两个连通块的最小候选边，依据切割性质可以安全选取；合并失败说明只会成环，跳过。

不变量是：任意时刻，并查集中的每个集合恰好对应一个已经内部连通、且所有内部连边都已计费的点集。 循环结束时所有点归入同一集合，cost 即为答案。

并查集用「路径压缩（这里是路径减半 parent[x] = parent[parent[x]]）+ 按大小合并」两个优化，让每次操作近乎常数时间。

解题步骤

构造虚拟井边：对 i 从 1 到 n，加入边 {0, i, wells[i-1]}。注意下标的错位——房子编号从 1 开始而 wells 数组从 0 开始，这里的 i - 1 写错就会整体错一位，代价全部张冠李戴。

并入管道边：pipes 的每一项 [u, v, w] 直接作为一条边。管道是无向的，Kruskal 天然按无向处理，不需要正反各存一条。

按权升序排序：Comparator.comparingInt(a -> a[2])。排序是 Kruskal 的前提，贪心「每次取当前最小的可用边」正确性来自切割性质。

并查集开 n + 1 个位置：编号 0 到 n。开成 n 会让虚拟节点或最后一座房子越界，这是虚拟节点法最典型的低级错误。

顺序扫边并合并：if (uf.union(e[0], e[1])) cost += e[2];。把「是否成功合并」直接作为「是否计费」的依据，是 Kruskal 最干净的写法——它同时完成了成环判定与去重（重复边的第二条必然合并失败），所以输入里的平行边不需要预处理。

find 用路径减半：parent[x] = parent[parent[x]] 在循环里就地把路径压扁一半，写法比递归的完全路径压缩短且没有栈开销。

union 按大小合并：把小树挂到大树下，避免退化成链。与路径压缩配合后单次操作近似 $O(\alpha)$。

累计恰好 n 次成功合并：新图有 n + 1 个点，生成树恰有 n 条边；达到该数量即可停止。虚拟井边保证新图一定连通。

以官方样例 n = 3、wells = [1, 2, 2]、pipes = [[1,2,1],[2,3,1]] 走一遍（答案 3）：

建图后的边集为：井边 (0,1,1)、(0,2,2)、(0,3,2)；管道边 (1,2,1)、(2,3,1)。

排序后（同权任意序）：(0,1,1)、(1,2,1)、(2,3,1)、(0,2,2)、(0,3,2)。

取 (0,1,1)：0 与 1 不同集合，合并成功，cost = 1。含义是在 1 号房子打井。
取 (1,2,1)：1 与 2 不同集合，合并成功，cost = 2。含义是从 1 号铺管道到 2 号。
取 (2,3,1)：合并成功，cost = 3、used = n，所有 n+1 个点已连通，直接停止；剩余两条井边若继续扫描都会成环。

返回 3，对应方案「在 1 号打井（1 元）+ 铺 1-2（1 元）+ 铺 2-3（1 元）」。

再看一个 pipes 为空的例子：n = 2、wells = [1, 1]、pipes = []。边集只有 (0,1,1) 与 (0,2,1)，两条都合并成功，cost = 2——每座房子各打一口井，符合直觉，主逻辑天然覆盖，不需要为「无管道」写任何特判。

最后看为什么不能先在原管道图里求 MST，再给每个连通块补最便宜的井：这种顺序会强迫原图的整个连通块都铺管。例如 wells = [1,2,2]、两条管道费用都为 5 时，它会花 10 + 1 = 11；最优却是三座房子分别打井，只花 5。虚拟节点把两类费用放进同一次排序，才能做全局权衡。

代码实现

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

class Solution {
    public int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
        List<int[]> edges = new ArrayList<>();
        // 打井 = 从虚拟水源 0 引一条边到该房子。
        for (int i = 1; i <= n; i++) {
            edges.add(new int[] {0, i, wells[i - 1]});
        }
        for (int[] p : pipes) {
            edges.add(new int[] {p[0], p[1], p[2]});
        }

        Collections.sort(edges, Comparator.comparingInt(a -> a[2]));

        UnionFind uf = new UnionFind(n + 1);
        int cost = 0;
        int used = 0;
        for (int[] e : edges) {
            // 合并成功才计费：失败说明两端已连通，这条边只会成环。
            if (uf.union(e[0], e[1])) {
                cost += e[2];
                if (++used == n) {
                    break;
                }
            }
        }
        return cost;
    }

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

        UnionFind(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);
            int 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;
        }
    }
}

import "sort"

type Edge struct {
	u int
	v int
	w int
}

type UnionFind struct {
	parent []int
	size   []int
}

func newUnionFind(n int) *UnionFind {
	parent := make([]int, n)
	size := make([]int, n)
	for i := 0; i < n; i++ {
		parent[i] = i
		size[i] = 1
	}
	return &UnionFind{parent: parent, size: size}
}

func (uf *UnionFind) find(x int) int {
	for x != uf.parent[x] {
		// 路径减半，就地把查找路径压扁。
		uf.parent[x] = uf.parent[uf.parent[x]]
		x = uf.parent[x]
	}
	return x
}

func (uf *UnionFind) union(a, b int) bool {
	ra := uf.find(a)
	rb := uf.find(b)
	if ra == rb {
		return false
	}
	// 按大小合并，小树挂到大树下。
	if uf.size[ra] < uf.size[rb] {
		ra, rb = rb, ra
	}
	uf.parent[rb] = ra
	uf.size[ra] += uf.size[rb]
	return true
}

func minCostToSupplyWater(n int, wells []int, pipes [][]int) int {
	edges := make([]Edge, 0, n+len(pipes))
	// 打井 = 从虚拟水源 0 引一条边到该房子。
	for i := 1; i <= n; i++ {
		edges = append(edges, Edge{u: 0, v: i, w: wells[i-1]})
	}
	for _, p := range pipes {
		edges = append(edges, Edge{u: p[0], v: p[1], w: p[2]})
	}

	sort.Slice(edges, func(i, j int) bool {
		return edges[i].w < edges[j].w
	})

	uf := newUnionFind(n + 1)
	cost := 0
	used := 0
	for _, e := range edges {
		// 合并成功才计费：失败说明两端已连通，这条边只会成环。
		if uf.union(e.u, e.v) {
			cost += e.w
			used++
			if used == n {
				break
			}
		}
	}
	return cost
}

复杂度分析

时间复杂度：$O((n + m) \log (n + m))$，其中 $m$ 为管道数。建边 $O(n + m)$，排序主导整体开销；随后扫一遍边，每条边做常数次并查集操作，配合路径压缩与按大小合并后单次近似 $O(\alpha(n))$（反阿克曼函数，可视为常数），这部分是 $O((n+m)\alpha(n))$。

空间复杂度：$O(n + m)$，边集合存下全部 $n + m$ 条边，并查集的 parent 与 size 各占 $O(n)$。排序若用归并还需 $O(n+m)$ 的辅助空间。

关键点总结

虚拟节点（超级源点）是把「点上的花费」翻译成「边上的花费」的标准手法。一旦两类决策被统一成同一种度量，原本要枚举的组合选择就自动被同一个贪心过程完成了。看到「每个连通块里至少要选一个点付费」就该想到它。

建模完成后要主动验证等价性——「合法方案 ↔ 连通子图」且代价相等，双向都说清楚，这是面试里区分「背过模板」和「真的会建模」的地方。

Kruskal 的「合并成功才计费」一行同时完成了成环判定、去重与计费，输入里的平行边、无序边都无需预处理。

虚拟节点会让点数变成 n + 1，并查集必须按新的点数开空间，编号 0 留给它——这是这类题最高频的越界来源。

房屋编号从 1 起而数组从 0 起，wells[i-1] 的错位要写对，写错不会崩溃只会静默给出错误答案，更难排查。

并查集的两个优化（路径压缩 + 按秩/大小合并）要一起用；只用其中一个在 $10^4$ 规模下也能过，但面试中应当写全并能说出为什么。

面试视角：本题可以直接答「加超级源点后求 MST」。边集显式且规模为 $n+m$，Kruskal 最顺手；若改成稠密图，通常比较邻接矩阵版 $O(V^2)$ Prim，避免枚举、排序海量边。

易错点总结

错误写法：并查集开 n 个位置而不是 n + 1。用例任意含 n 号房子的输入：访问 parent[n] 直接数组越界抛异常。

错误写法：井边写成 {0, i, wells[i]}。用例 n = 3, wells = [1,2,2]：每座房子拿到的是下一座房子的打井费，且 wells[3] 越界；即使侥幸不越界，代价也整体错位一格。

错误写法：忘记把井边和管道边放进同一个排序集合，而是先对管道求 MST、再给每个连通块加一口最便宜的井。用例 n = 3, wells = [1,2,2], pipes = [[1,2,5],[2,3,5]]：这种写法会先花 10 元铺两条管道再加 1 元打井共 11 元，而正确答案是三口井各打共 5 元——两类花费不放在一起比较就无法做全局权衡。

错误写法：无论 union 是否成功都累加边权。用例 n = 3, wells = [1,2,2], pipes = [[1,2,1],[2,3,1]]：五条边全部计费得 7，而正确答案是 3，多付的正是会成环的冗余边。

错误写法：判完 find(a) != find(b) 后却直接写 parent[a] = b。若 a 不是根，这会把它的子树从原集合拆开，后续连通性判断失真；必须连接两个根节点。

错误写法：忘记排序直接按输入顺序扫边。用例 n = 3, wells = [1,2,2], pipes = [[1,2,1],[2,3,1]]：先取到的可能是权为 2 的井边，贪心失去依据，结果偏大。

错误写法：提前停止条件写成成功合并 n - 1 次。新图共有 n + 1 个点，MST 需要 n 条边；少一次合并会留下一个房子未接通水源。

相似题目

题目	难度	考察点
1135. 最低成本连通所有城市	中等	最纯粹的 Kruskal 模板，没有虚拟节点，且需判断图是否本就不连通
1584. 连接所有点的最小费用	中等	边需要由坐标两两生成（稠密图），更适合 Prim，与本题的选型形成对照
1489. 找到最小生成树里的关键边和伪关键边	困难	在 MST 之上追问每条边的必需性，要反复删边/强制加边重跑 Kruskal
684. 冗余连接	中等	只用并查集的成环判定，是本题「合并失败即成环」这一步的单独练习
547. 省份数量	中等	纯连通块计数，帮助建立「并查集集合即连通块」的直觉
778. 水位上升的泳池中游泳	困难	按边权升序并查集合并直到两点连通，是 Kruskal 思想的另一种变形