统计可能的树根数目

题目

2581. 统计可能的树根数目

---

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 <= 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

题解

方法一：

思路

我们可以尝试暴力的做法，对于每个节点都尝试作为根，然后从该节点开始深搜或广搜统计bob的猜想有多少个为true，即为cnt，这样做的时间复杂度$O(n^2)$。

我们发现如果已经计算了以x为根的cnt。
那么与x相邻的y其实是可以利用x的统计数来优化的。
试想从x作为根到y作为根其实除了x到y这条边之外其余边都不需要发生变化。
因此可以$O(1)$计算出相邻节点的cnt。

我们可以先通过dfs求出任意节点的cnt，然后再从该节点开始dfs求出其他节点的cnt。对于每个节点的cnt若大于等于k则增加统计数。

代码

class Solution {
public:
    int cnt = 0, ans = 0, K = 0;
    set<pair<int,int>> st;
    vector<int> g[100005];
    void getcnt(int u, int fa) {
        if (st.count({fa,u})) {
            cnt++;
        }
        for (int i:g[u]) {
            if (fa == i) continue;
            getcnt(i, u);
        }
    }
    void dfs(int u, int fa) {
        if (st.count({fa,u})) {
            cnt--;
        }
        if (st.count({u, fa})) {
            cnt++;
        }
        
        // cout << u << " " << fa << " " << cnt << " " << (cnt>=K) << endl;
        ans += cnt >= K;
        for (int i:g[u]) {
            if (fa == i) continue;
            dfs(i, u);
        }
        if (st.count({fa,u})) {
            cnt++;
        }
        if (st.count({u, fa})) {
            cnt--;
        }
    }
    
    int rootCount(vector<vector<int>>& edges, vector<vector<int>>& gu, int k) {
        K = k;
        for (auto& i:edges) {
            g[i[0]].push_back(i[1]);
            g[i[1]].push_back(i[0]);
        }
        for (auto& i:gu) {
            st.insert({i[0], i[1]});
        }
        getcnt(0, -1);
        dfs(0, -1);
        return ans;
    }
};