leetcode 2581. 统计可能的树根数目【换根dp】

原题链接：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

解题思路：

只是一个比较简单的困难题，我们需要求的是可能的根节点有多少个，也就是说最暴力的做法就是枚举每一个点作为根节点，然后考虑这种情况下，猜测中有多少个猜测满足要求，如果满足要求的猜测个数>=k,那么说明这种情况这个点可以作为根节点，但是这个题目n=1e5,如果直接暴力枚举每一个根结点，时间复杂度为O(n^2),这个时间复杂度肯定过不了，所以这个题目肯定不能直接枚举每一个根结点，考虑优化，需要枚举每一个根节点，也就是需要换根，我们可以换根dp进行优化，换根dp实际上就是俩次dfs,第一次dfs先确定某一个结点为根结点，计算这种情况会有多少个猜测成立，第二次dfs就是考虑换根的影响，我们不妨画个图描述一下：

如上图所示，左边图最开始0号结点为根结点，右图所示为换根操作，把根换为1号点，我们可以发现换根之后相对于换根的前一步比较，只有换根的俩个点之间父子关系发生翻转，另外的其他所有点之间的父子关系不变，发现这个性质之后我们就可以根据换根造成的影响来计算这个新树中满足要求的猜测数量，然后判断这个数量是否大于等于k，从而更新答案。

时间复杂度：O(n+m)，n表示点数，m表示guesses的长度。

空间复杂度：O(n+m)，n表示点数，m表示guesses的长度。

cpp代码如下:

class Solution {
public:
    int rootCount(vector<vector<int>>& edges, vector<vector<int>>& guesses, int k) {
        int ans=0,n=edges.size()+1;
        vector<vector<int>>g(n,vector<int>());
        for(auto& t:edges){   //邻接表建图
            int x=t[0],y=t[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }

        map<pair<int,int>,int>mp;  //首先把所有的猜测存起来，用于后续判断
        for(auto& t:guesses){
            int x=t[0],y=t[1];
            mp[pair{x,y}]++;
        }

        function<int(int,int)>dfs1=[&](int x,int fa)->int {
            int res=0;
            for(auto& y:g[x]){
                if(y==fa)continue;
                if(mp.count(pair{x,y})){
                    res++;
                }
                res+=dfs1(y,x);
            }
            return res;
        };
        //首先考虑0号结点为根结点时，会有多少个猜测满足要求，用cnt记录
        int cnt=dfs1(0,-1);
        if(cnt>=k)ans++;  //这种情况如果满足要求，更新答案

        function<void(int,int,int)>dfs2=[&](int x,int fa,int cnt){
            for(auto& y:g[x]){
                if(y==fa)continue;
                int val=cnt;  //每个子树都是基于cnt进行的，所以先用一个val存起来
                //换根相当于把x->y的边变为y->x的边，其他的边不变
                if(mp.count(pair{x,y})){  //把x->y的边减去
                    val--;
                }
                if(mp.count(pair{y,x})){
                    val++;   //把y->x的加上
                }

                if(val>=k)ans++;   //更新答案
                dfs2(y,x,val);
            }
        };

        //然后考虑换根操作
        dfs2(0,-1,cnt);
        return ans;
    }
};