2581. 统计可能的树根数目

来源第 99 场双周赛 Q4难度困难分数2228

题目描述

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

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

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

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

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 <= 10⁵
1 <= guesses.length <= 10⁵
0 <= a_i, b_i, u_j, v_j <= n - 1
a_i != b_i
u_j != v_j
edges 表示一棵有效的树。
guesses[j] 是树中的一条边。
guesses 是唯一的。
0 <= k <= guesses.length

解法

方法一：树形 DP（换根）

思考

树的根未知，猜测是有向边。若以某点为根时猜中不少于 \(k\) 条，则该点可行。对每个根单独 DFS 需 \(O(n^2)\)。

先以 \(0\) 为根统计猜中数 \(\textit{cnt}\)。换根到邻点 \(j\) 时，边 \((i,j)\) 不再由父指向子， \((j,i)\) 则变为由父指向子，故 \(\textit{cnt}\) 增减对应猜测是否存在。第二次 DFS 遍历每个根并计数。

我们先遍历题目给定的边集合 \(\textit{edges}\)，将其转换为邻接表 \(g\)，其中 \(g[i]\) 表示节点 \(i\) 的所有邻接节点。用哈希表 \(gs\) 记录题目给定的猜测集合 \(\textit{guesses}\)。

接下来，我们先从节点 \(0\) 开始，进行一次 DFS，统计从节点 \(0\) 出发，能够到达的所有节点中，有多少条边在 \(\textit{guesses}\) 中。我们用变量 \(\textit{cnt}\) 记录这个数量。

然后，我们再从节点 \(0\) 开始，进行一次 DFS，统计以每个点为根的树中，有多少条边在 \(\textit{guesses}\) 中。如果这个数量大于等于 \(k\)，则说明该节点是一个可能的根节点，我们将答案加 \(1\)。

因此，问题就转化为了求以每个点为根的树中，有多少条边在 \(\textit{guesses}\) 中。我们已经知道了从节点 \(0\) 出发，能够到达的所有节点中，有多少条边在 \(\textit{guesses}\) 中，即 \(\textit{cnt}\)。我们可以通过在 DFS 中，将 \(\textit{cnt}\) 的值加上或减去当前边是否在 \(\textit{guesses}\) 中，来维护这个值。

假设我们当前遍历到节点 \(i\)，此时 \(\textit{cnt}\) 表示以 \(i\) 为根节点，有多少条边在 \(\textit{guesses}\) 中。那么对于 \(i\) 的每个邻接节点 \(j\)，我们要计算以 \(j\) 为根节点，有多少条边在 \(\textit{guesses}\) 中。如果 \((i, j)\) 在 \(\textit{guesses}\) 中，那么以 \(j\) 为根节点的，就不存在 \((i, j)\) 这条边，因此 \(\textit{cnt}\) 的值要减去 \(1\)。如果 \((j, i)\) 在 \(\textit{guesses}\) 中，那么以 \(j\) 为根节点的，要多出一条 \((i, j)\) 这条边，因此 \(\textit{cnt}\) 的值要加上 \(1\)。即 \(f[j] = f[i] + (j, i) \in \textit{guesses} - (i, j) \in \textit{guesses}\)。其中 \(f[i]\) 表示以 \(i\) 为根节点，有多少条边在 \(\textit{guesses}\) 中。

时间复杂度 \(O(n + m)\)，空间复杂度 \(O(n + m)\)。其中 \(n\) 和 \(m\) 分别为 \(\textit{edges}\) 和 \(\textit{guesses}\) 的长度。

相似题目：

834. 树中距离之和

Python3JavaC++GoTypeScript

class Solution:
    def rootCount(
        self, edges: List[List[int]], guesses: List[List[int]], k: int
    ) -> int:
        def dfs1(i, fa):
            nonlocal cnt
            for j in g[i]:
                if j != fa:
                    cnt += gs[(i, j)]
                    dfs1(j, i)

        def dfs2(i, fa):
            nonlocal ans, cnt
            ans += cnt >= k
            for j in g[i]:
                if j != fa:
                    cnt -= gs[(i, j)]
                    cnt += gs[(j, i)]
                    dfs2(j, i)
                    cnt -= gs[(j, i)]
                    cnt += gs[(i, j)]

        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)
        gs = Counter((u, v) for u, v in guesses)
        cnt = 0
        dfs1(0, -1)
        ans = 0
        dfs2(0, -1)
        return ans

class Solution {
    private List<Integer>[] g;
    private Map<Long, Integer> gs = new HashMap<>();
    private int ans;
    private int k;
    private int cnt;
    private int n;

    public int rootCount(int[][] edges, int[][] guesses, int k) {
        this.k = k;
        n = edges.length + 1;
        g = new List[n];
        Arrays.setAll(g, e -> new ArrayList<>());
        for (var e : edges) {
            int a = e[0], b = e[1];
            g[a].add(b);
            g[b].add(a);
        }
        for (var e : guesses) {
            int a = e[0], b = e[1];
            gs.merge(f(a, b), 1, Integer::sum);
        }
        dfs1(0, -1);
        dfs2(0, -1);
        return ans;
    }

    private void dfs1(int i, int fa) {
        for (int j : g[i]) {
            if (j != fa) {
                cnt += gs.getOrDefault(f(i, j), 0);
                dfs1(j, i);
            }
        }
    }

    private void dfs2(int i, int fa) {
        ans += cnt >= k ? 1 : 0;
        for (int j : g[i]) {
            if (j != fa) {
                int a = gs.getOrDefault(f(i, j), 0);
                int b = gs.getOrDefault(f(j, i), 0);
                cnt -= a;
                cnt += b;
                dfs2(j, i);
                cnt -= b;
                cnt += a;
            }
        }
    }

    private long f(int i, int j) {
        return 1L * i * n + j;
    }
}

class Solution {
public:
    int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses, int k) {
        int n = edges.size() + 1;
        vector<int> g[n];
        unordered_map<long long, int> gs;
        auto f = [&](int i, int j) {
            return 1LL * i * n + j;
        };
        for (auto& e : edges) {
            int a = e[0], b = e[1];
            g[a].push_back(b);
            g[b].push_back(a);
        }
        for (auto& e : guesses) {
            int a = e[0], b = e[1];
            gs[f(a, b)]++;
        }
        int ans = 0;
        int cnt = 0;

        function<void(int, int)> dfs1 = [&](int i, int fa) {
            for (int& j : g[i]) {
                if (j != fa) {
                    cnt += gs[f(i, j)];
                    dfs1(j, i);
                }
            }
        };

        function<void(int, int)> dfs2 = [&](int i, int fa) {
            ans += cnt >= k;
            for (int& j : g[i]) {
                if (j != fa) {
                    int a = gs[f(i, j)];
                    int b = gs[f(j, i)];
                    cnt -= a;
                    cnt += b;
                    dfs2(j, i);
                    cnt -= b;
                    cnt += a;
                }
            }
        };
        dfs1(0, -1);
        dfs2(0, -1);
        return ans;
    }
};

func rootCount(edges [][]int, guesses [][]int, k int) (ans int) {
    n := len(edges) + 1
    g := make([][]int, n)
    gs := map[int]int{}
    for _, e := range edges {
        a, b := e[0], e[1]
        g[a] = append(g[a], b)
        g[b] = append(g[b], a)
    }
    f := func(i, j int) int {
        return i*n + j
    }
    for _, e := range guesses {
        a, b := e[0], e[1]
        gs[f(a, b)]++
    }

    cnt := 0
    var dfs1 func(i, fa int)
    var dfs2 func(i, fa int)
    dfs1 = func(i, fa int) {
        for _, j := range g[i] {
            if j != fa {
                cnt += gs[f(i, j)]
                dfs1(j, i)
            }
        }
    }
    dfs2 = func(i, fa int) {
        if cnt >= k {
            ans++
        }
        for _, j := range g[i] {
            if j != fa {
                a, b := gs[f(i, j)], gs[f(j, i)]
                cnt -= a
                cnt += b
                dfs2(j, i)
                cnt -= b
                cnt += a
            }
        }
    }
    dfs1(0, -1)
    dfs2(0, -1)
    return
}

function rootCount(edges: number[][], guesses: number[][], k: number): number {
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n }, () => []);
    const gs: Map<number, number> = new Map();
    const f = (i: number, j: number) => i * n + j;
    for (const [a, b] of edges) {
        g[a].push(b);
        g[b].push(a);
    }
    for (const [a, b] of guesses) {
        const x = f(a, b);
        gs.set(x, gs.has(x) ? gs.get(x)! + 1 : 1);
    }
    let ans = 0;
    let cnt = 0;
    const dfs1 = (i: number, fa: number): void => {
        for (const j of g[i]) {
            if (j !== fa) {
                cnt += gs.get(f(i, j)) || 0;
                dfs1(j, i);
            }
        }
    };
    const dfs2 = (i: number, fa: number): void => {
        ans += cnt >= k ? 1 : 0;
        for (const j of g[i]) {
            if (j !== fa) {
                const a = gs.get(f(i, j)) || 0;
                const b = gs.get(f(j, i)) || 0;
                cnt -= a;
                cnt += b;
                dfs2(j, i);
                cnt -= b;
                cnt += a;
            }
        }
    };
    dfs1(0, -1);
    dfs2(0, -1);
    return ans;
}

方法二：换根 DP + 显式栈

思考

树的根未知，每条猜测是一条有向边。以某点为根时，猜中条数不少于 \(k\) 才算可行。 \(n\le 10^5\)，对每个根各遍历一次是平方级；沿链递归时调用深度等于节点数，链长达到 \(1000\) 就会超出 Python 的递归上限。猜中数在换根时只变一条边：从父亲 \(i\) 走到孩子 \(j\)， \((i,j)\) 不再由父指向子， \((j,i)\) 变为由父指向子。第一遍用栈从 \(0\) 累加方向与父子关系一致的猜测，得到 \(\textit{cnt}\)。第二遍把调整后的计数放进栈帧，弹出时若计数不少于 \(k\) 就记一个可行根，再把 \(\textit{cnt}-(i,j)+(j,i)\) 传给孩子。

先把 \(\textit{edges}\) 转成邻接表 \(g\)，并用哈希表 \(gs\) 记录 \(\textit{guesses}\)。用显式栈从节点 \(0\) 遍历整棵树。走到孩子 \(j\) 时，若 \((i,j)\) 在 \(gs\) 中，就把 \(\textit{cnt}\) 加上对应次数。遍历结束后， \(\textit{cnt}\) 就是以 \(0\) 为根时猜中的条数。

再换根。栈里存放节点、父亲和以该节点为根时的猜中数。初始帧是 \((0,-1,\textit{cnt})\)。弹出节点 \(i\) 时，若当前计数不少于 \(k\)，答案加 \(1\)。对于邻居 \(j\)，以 \(j\) 为根时 \((i,j)\) 不再由父指向子， \((j,i)\) 变为由父指向子，所以传给孩子的计数是 \(c-gs[(i,j)]+gs[(j,i)]\)。