小美的修路

美团笔试编程题 · 带必选边的最小生成树

题目描述

地图上有 n 个城市，小美准备修建一些道路，使得任意两个城市之间都能通过道路 到达。现在给定 m 个修路计划（部分计划是必选的），请你帮小美规划出最优方案： 在保证连通且包含全部必选计划的前提下，让总花费最小。

输入格式

第一行两个正整数 n、m，用空格隔开，代表城市数量与计划数量。 接下来 m 行，每行四个正整数 u、v、w、required： 在城市 u 和 v 之间修建一条花费为 w 的道路；required 为 1 表示必选，0 表示可选。

输出格式

如果无解（无法使所有城市连通），输出 -1。 否则第一行输出一个正整数，代表选择的计划数量；第二行输出这些计划的编号 （1 起），用空格隔开。

示例

示例 1

输入：

3 4
1 2 3 1
1 2 2 0
1 3 1 0
2 3 3 0

输出：

2
1 3

说明： 方案 1 是必选的（花费 3）；方案 2、3 之间要连通城市 3， 只能选更便宜的方案 3（花费 1）。总花费 3 + 1 = 4。 若选方案 1、2，总花费 5 更大。

---

解题思路

第一步：先看清「带必选边」的约束

这是一道最小生成树（MST）题，但多了两个条件：

• 有些边必须选（哪怕它很贵）；
• 要输出具体的选边集合，而不只是总花费。

第二步：普通 MST 的做法

Kruskal 算法：把所有边按花费升序排序，用并查集依次尝试加入， 能连通两个不同集合的就选，直到连通所有点。贪心的正确性由「切割性质」保证。

第三步：加入必选边

必选边无论如何都要选，所以把它们先无条件加入并查集：

for u, v, w, req, i in edges:
    if req == 1:
        union(u, v)
        chosen.append(i)

注意必选边之间即使成环也照单全收（题目要求必须包含它们）。 并查集在这里只负责回答「还有哪些点没连通」。

第四步：补齐剩下的连通性

把可选边按花费升序，继续 Kruskal：

for u, v, w, req, i in sorted(edges, key=lambda e: e[2]):
    if req == 0 and union(u, v):
        chosen.append(i)

第五步：判断无解

所有边都试过之后，如果还有城市和城市 1 不在同一集合，说明无论如何 都连通不了，输出 -1。

---

完整代码实现

算法可视化 · 逐帧回放这道题的代码115 帧 · 录制输入

import sys


def solve() -> None:
    """并查集：必选边无条件加入，可选边按花费贪心"""
    data = sys.stdin.read().split()
    if not data:
        return
    idx = 0
    n = int(data[idx]); m = int(data[idx + 1]); idx += 2

    edges = []
    for i in range(1, m + 1):
        u = int(data[idx]); v = int(data[idx + 1])
        w = int(data[idx + 2]); req = int(data[idx + 3])
        idx += 4
        edges.append((u, v, w, req, i))

    parent = list(range(n + 1))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]   # 路径压缩
            x = parent[x]
        return x

    def union(x: int, y: int) -> bool:
        rx, ry = find(x), find(y)
        if rx == ry:
            return False
        parent[rx] = ry
        return True

    chosen = []
    # 1. 必选边全部加入
    for u, v, w, req, i in edges:
        if req == 1:
            union(u, v)
            chosen.append(i)

    # 2. 可选边按花费升序，Kruskal
    for u, v, w, req, i in sorted(edges, key=lambda e: e[2]):
        if req == 0 and union(u, v):
            chosen.append(i)

    # 3. 判断是否连通
    root = find(1)
    for city in range(2, n + 1):
        if find(city) != root:
            sys.stdout.write("-1\n")
            return

    chosen.sort()
    sys.stdout.write(str(len(chosen)) + "\n" + " ".join(map(str, chosen)) + "\n")


if __name__ == "__main__":
    solve()

---

示例推演

输入 4 个计划：

编号	边	花费	必选
1	1-2	3	是
2	1-2	2	否
3	1-3	1	否
4	2-3	3	否

1. 必选边：加入 1（1-2 连通），已选 {1}；
2. 可选边按花费升序：1（方案 3）、2（方案 2）、3（方案 4、方案 1）：
   • 方案 3（1-3，花费 1）：1、3 未连通 → 选，已选 {1, 3}；
   • 方案 2（1-2，花费 2）：已连通 → 跳过；
   • 方案 4（2-3，花费 3）：已连通 → 跳过；
3. 所有城市连通，输出数量 2 与编号 1 3。

---

复杂度分析

解法	时间复杂度	空间复杂度	说明
并查集 + Kruskal	O(m log m)	O(n + m)	推荐（瓶颈在排序）

---

易错点总结

1. 必选边没有全部加入

必选边要无条件加入，即使它与已有边成环。漏加会改变连通性与答案。

2. 忘记输出方案编号

本题不止要最小花费，还要输出选中的计划编号（1 起）。 只算花费不看编号，就会得到「结果对但被判错」。

3. 无解时只输出 -1

无解时不要输出数量那一行，直接输出 -1 即可。

4. 可选边排序的正确性

Kruskal 要求按花费升序处理可选边；如果混入必选边一起排序， 会因为「必选边已经被加入」而判重，虽不影响正确性但代码更绕。

---

扩展思考

• 若要求「选出的道路数量最少」，需要在最小花费的基础上再比数量。
• 若必选边之间成环，允许保留（本题如此）；若题目要求「生成树」 且禁止成环，则必选边成环时应判无解。

---

相关题目

• 小美的彩带