2581-统计可能的树根数目

Raphael Liu Lv10
  2023-09-13 00:47:26 2023-09-13 00:47:26 Created   2023-09-13 16:06:29 2023-09-13 16:06:29 Updated    

Alice 有一棵 n 个节点的树，节点编号为 0 到 n - 1 。树用一个长度为 n - 1 的二维整数数组 edges 表示，其中
edges[i] = [ai, bi] ，表示树中节点 ai 和 bi 之间有一条边。

Alice 想要 Bob 找到这棵树的根。她允许 Bob 对这棵树进行若干次 猜测 。每一次猜测，Bob 做如下事情：

  • 选择两个 不相等 的整数 u 和 v ，且树中必须存在边 [u, v] 。
  • Bob 猜测树中 u 是 v 的 父节点 。

Bob 的猜测用二维整数数组 guesses 表示，其中 guesses[j] = [uj, vj] 表示 Bob 猜 uj 是 vj
的父节点。

Alice 非常懒，她不想逐个回答 Bob 的猜测，只告诉 Bob 这些猜测里面 至少 有 k 个猜测的结果为 true 。

给你二维整数数组 edges ，Bob 的所有猜测和整数 k ，请你返回可能成为树根的 节点数目 。如果没有这样的树，则返回 0。

示例 1：

**输入：** edges = [[0,1],[1,2],[1,3],[4,2]], guesses = [[1,3],[0,1],[1,0],[2,4]], k = 3
**输出：** 3
**解释：**
根为节点 0 ，正确的猜测为 [1,3], [0,1], [2,4]
根为节点 1 ，正确的猜测为 [1,3], [1,0], [2,4]
根为节点 2 ，正确的猜测为 [1,3], [1,0], [2,4]
根为节点 3 ，正确的猜测为 [1,0], [2,4]
根为节点 4 ，正确的猜测为 [1,3], [1,0]
节点 0 ，1 或 2 为根时，可以得到 3 个正确的猜测。

示例 2：

**输入：** edges = [[0,1],[1,2],[2,3],[3,4]], guesses = [[1,0],[3,4],[2,1],[3,2]], k = 1
**输出：** 5
**解释：**
根为节点 0 ，正确的猜测为 [3,4]
根为节点 1 ，正确的猜测为 [1,0], [3,4]
根为节点 2 ，正确的猜测为 [1,0], [2,1], [3,4]
根为节点 3 ，正确的猜测为 [1,0], [2,1], [3,2], [3,4]
根为节点 4 ，正确的猜测为 [1,0], [2,1], [3,2]
任何节点为根，都至少有 1 个正确的猜测。

提示：

  • edges.length == n - 1
  • 2 <= n <= 105
  • 1 <= guesses.length <= 105
  • 0 <= ai, bi, uj, vj <= n - 1
  • ai != bi
  • uj != vj
  • edges 表示一棵有效的树。
  • guesses[j] 是树中的一条边。
  • guesses 是唯一的。
  • 0 <= k <= guesses.length

如果只求以 0 为根时的猜对次数 cnt}_0，那么把 guesses 转成哈希表，DFS 一次这棵树就可以算出来。

如果要枚举以每个点为根时的猜对次数，暴力做法就太慢了，怎么优化呢？

注意到，如果节点 x 和 y 之间有边，那么从「以 x 为根的树」变成「以 y 为根的树」，就只有 [x,y] 和 [y,x] 这两个猜测的正确性变了，其余猜测的正确性不变。

因此，从 0 出发，再次 DFS 这棵树，从节点 x 递归到节点 y 时：

  • 如果有猜测 [x,y]，那么猜对次数减一；
  • 如果有猜测 [y,x]，那么猜对次数加一。

DFS 的同时，统计猜对次数 \ge k 的节点个数，即为答案。

这个套路叫做「换根 DP」。

附：视频讲解

[sol1-Python3]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution:
    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:
        g = [[] for _ in range(len(edges) + 1)]
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)  # 建图

        s = {(x, y) for x, y in guesses}  # guesses 转成哈希表 s

        ans = cnt0 = 0
        def dfs(x: int, fa: int) -> None:
            nonlocal cnt0
            for y in g[x]:
                if y != fa:
                    cnt0 += (x, y) in s  # 以 0 为根时，猜对了
                    dfs(y, x)
        dfs(0, -1)

        def reroot(x: int, fa: int, cnt: int) -> None:
            nonlocal ans
            ans += cnt >= k  # 此时 cnt 就是以 x 为根时的猜对次数
            for y in g[x]:
                if y != fa:
                    reroot(y, x, cnt - ((x, y) in s) + ((y, x) in s))
        reroot(0, -1, cnt0)
        return ans
[sol1-Java]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class Solution {
    private List<Integer>[] g;
    private Set<Long> s = new HashSet<>();
    private int k, ans, cnt0;

    public int rootCount(int[][] edges, int[][] guesses, int k) {
        this.k = k;
        g = new ArrayList[edges.length + 1];
        Arrays.setAll(g, e -> new ArrayList<>());
        for (var e : edges) {
            int x = e[0], y = e[1];
            g[x].add(y);
            g[y].add(x); // 建图
        }

        for (var e : guesses) // guesses 转成哈希表
            s.add((long) e[0] << 32 | e[1]); // 两个 4 字节数压缩成一个 8 字节数

        dfs(0, -1);
        reroot(0, -1, cnt0);
        return ans;
    }

    private void dfs(int x, int fa) {
        for (var y : g[x])
            if (y != fa) {
                if (s.contains((long) x << 32 | y)) // 以 0 为根时，猜对了
                    ++cnt0;
                dfs(y, x);
            }
    }

    private void reroot(int x, int fa, int cnt) {
        if (cnt >= k) ++ans; // 此时 cnt 就是以 x 为根时的猜对次数
        for (var y : g[x])
            if (y != fa) {
                int c = cnt;
                if (s.contains((long) x << 32 | y)) --c; // 原来是对的，现在错了
                if (s.contains((long) y << 32 | x)) ++c; // 原来是错的，现在对了
                reroot(y, x, c);
            }
    }
}
[sol1-C++]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
public:
    int rootCount(vector<vector<int>> &edges, vector<vector<int>> &guesses, int k) {
        vector<vector<int>> g(edges.size() + 1);
        for (auto &e : edges) {
            int x = e[0], y = e[1];
            g[x].push_back(y);
            g[y].push_back(x); // 建图
        }

        unordered_set<long> s;
        for (auto &e : guesses) // guesses 转成哈希表
            s.insert((long) e[0] << 32 | e[1]); // 两个 4 字节数压缩成一个 8 字节数

        int ans = 0, cnt0 = 0;
        function<void(int, int)> dfs = [&](int x, int fa) {
            for (int y : g[x])
                if (y != fa) {
                    cnt0 += s.count((long) x << 32 | y); // 以 0 为根时，猜对了
                    dfs(y, x);
                }
        };
        dfs(0, -1);

        function<void(int, int, int)> reroot = [&](int x, int fa, int cnt) {
            ans += cnt >= k; // 此时 cnt 就是以 x 为根时的猜对次数
            for (int y : g[x])
                if (y != fa) {
                    reroot(y, x, cnt
                                 - s.count((long) x << 32 | y) // 原来是对的，现在错了
                                 + s.count((long) y << 32 | x)); // 原来是错的，现在对了
                }
        };
        reroot(0, -1, cnt0);
        return ans;
    }
};
[sol1-Go]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
func rootCount(edges [][]int, guesses [][]int, k int) (ans int) {
	g := make([][]int, len(edges)+1)
	for _, e := range edges {
		v, w := e[0], e[1]
		g[v] = append(g[v], w)
		g[w] = append(g[w], v) // 建图
	}

	type pair struct{ x, y int }
	s := make(map[pair]int, len(guesses))
	for _, p := range guesses { // guesses 转成哈希表
		s[pair{p[0], p[1]}] = 1
	}

	cnt0 := 0
	var dfs func(int, int)
	dfs = func(x, fa int) {
		for _, y := range g[x] {
			if y != fa {
				if s[pair{x, y}] == 1 { // 以 0 为根时，猜对了
					cnt0++
				}
				dfs(y, x)
			}
		}
	}
	dfs(0, -1)

	var reroot func(int, int, int)
	reroot = func(x, fa, cnt int) {
		if cnt >= k { // 此时 cnt 就是以 x 为根时的猜对次数
			ans++
		}
		for _, y := range g[x] {
			if y != fa {
				reroot(y, x, cnt-s[pair{x, y}]+s[pair{y, x}])
			}
		}
	}
	reroot(0, -1, cnt0)
	return
}

复杂度分析

  • 时间复杂度：O(n+m)，其中 n 为 edges 的长度加一，m 为 guesses 的长度。
  • 空间复杂度：O(n+m)。

相似题目

更多题目，可以看我的 算法竞赛模板库 中的「换根 DP」。

思考题

如果把「u 是 v 的父节点」改成「u 是 v 的祖先节点」，要怎么做呢？

如果改成「guesses}[i] 猜对会得到 score}[i] 分，计算的是以每个点为根时的得分之和」，要怎么做呢？（本题相当于 score}[i] 均为 1）

 Comments
On this page
2581-统计可能的树根数目
  1. 复杂度分析
  2. 相似题目
  3. 思考题