2581. 统计可能的树根数目(2228)

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：

1
2
3
4
5
6
7
8
9
输入：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：

1
2
3
4
5
6
7
8
9
输入：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 <= 10^5
  • 1 <= guesses.length <= 10^5
  • 0 <= ai, bi, uj, vj <= n - 1
  • ai != bi
  • uj != vj
  • edges 表示一棵有效的树。
  • guesses[j] 是树中的一条边。
  • guesses 是唯一的。
  • 0 <= k <= guesses.length

思路：换根操作在reroot函数中，在计算出 zeroRootRes 后，我们可以再次从 0 出发，DFS 这棵树。从节点 x 递归到节点 y 时：

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

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
44
class Solution {
    public int rootCount(int[][] edges, int[][] guesses, int k) {
        int n = edges.length+1, res = 0;
        HashSet<Integer> set = new HashSet<>();
        for (int guess[] : guesses) {
            set.add(guess[0]*n + guess[1]);
        }
        List<Integer> g[] = new ArrayList[n];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (int edge[] : edges) {
            int x = edge[0], y = edge[1];
            g[x].add(y);
            g[y].add(x);
        }
        int zeroRootRes = dfs(g, set, -1, 0);  // 计算以0为root节点的猜测结果
        int cnt[] = new int[n];
        reroot(g, set, -1, 0, cnt);  // 计算以其他节点为root节点的猜测结果
        for (int i = 0; i < n; i++) {
            if (zeroRootRes + cnt[i] >= k) res++;
        }
        return res;
    }
    public int dfs(List<Integer> g[], HashSet<Integer> set, int pre, int cur) {
        int res = 0, n = g.length;
        if (pre != -1 && set.contains(pre*n + cur)) res++;
        for (int nxt : g[cur]) {
            if (nxt == pre) continue;
            res += dfs(g, set, cur, nxt);
        }
        return res;
    }
    public void reroot(List<Integer> g[], HashSet<Integer> set, int pre, int cur, int cnt[]) {
        int n = g.length;
        if (pre != -1) {
            cnt[cur] = cnt[pre];
            if (set.contains(pre*n + cur)) cnt[cur]--;
            if (set.contains(cur*n + pre)) cnt[cur]++;
        }
        for (int nxt : g[cur]) {
            if (nxt == pre) continue;
            reroot(g, set, cur, nxt, cnt);
        }
    }
}