LeetCode每日一题【c++版】- leetcode 2581. 统计可能的树根数目【换根dp】

题目描述

题目链接：2581. 统计可能的树根数目

输入输出描述

示例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 个正确的猜测。

解题思路

        通常来说这种题目的套路就是先随便定一个根，比如0作为整棵树的根节点，然后从0开始进行DFS，计算出其他所有子树i中，属于guesses中的边有多少条，记录在f[i]中，这一步很简单，就是个简单的树形DP。

        然后再进行一次DFS进行状态转移，关键就在于这个状态转移的递归函数怎么实现。dfs(u,fa,top)表示考虑以u作为整棵树的根，top表示把根调整为u之前（即以u的父节点fa为根时）属于u上面（祖先方向）的子图中有多少条猜测中的边。此时，如果递归到u的某个孩子v时，我们就要计算v的top是多少;

    gain=f[u]−f[v]就是v的其他兄弟节点对v节点top值的额外贡献，因为其他兄弟节点在v换成根之后会成为它的子树。
    如果u->v（箭头表示父节点指向子节点）这条边在guesses中，就还需要在gain中减去这条边。但如果v->u这条边在guesses中，gain就要加上这条边，此时我们得到v的top值更新为top=top+gain

        每当遍历到一个节点u，只要满足f[u]+top>=k,说明这个节点作为整棵树的根时可以保证guesses中至少有k条正确，答案ans自增即可。第二次DFS遍历完成后，返回ans的值即可。两次dfs的说明如下：

1. 第一次找到以某个节点（比如 0）作为根节点时的正确数 zero_root_correct，作为状态转移的初始值。
2. 第二次我们通过遍历到父子节点，来对正确数进行转移，从而得到每个节点作为根节点时的正确数 correct。

复杂度

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

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

代码

#include <iostream>
#include <vector>
#include <unordered_set>
#include <set>
using namespace std;
class Solution {
public:
    int rootCount(vector<vector<int>> &edges, vector<vector<int>> &guesses, int k)
    {
        int n = edges.size() + 1;  // n个节点
        // 获取每个节点的邻居节点
        vector<vector<int>> neighbors(n);
        for (auto &e : edges) {
            int x = e[0], y = e[1];
            neighbors[x].emplace_back(y);
            neighbors[y].emplace_back(x);
        }
        // 生成猜测数对的id = u * K + v
        unordered_set<long long> guessesIds;
        for (auto &g : guesses) {
            int u = g[0], v = g[1];
            guessesIds.emplace(u * K + v);
        }
        // 获取以0节点为根时的正确数
        int zeroRootCorrect = dfs_GetZeroRootCorrect(0, -1, neighbors, guessesIds);
        // 以0节点为根，其正确数为起始值，再一次遍历所有节点，并得到其他节点作为根节点的正确数
        dfs_ChangeRoot(0, -1, k, zeroRootCorrect, neighbors, guessesIds);
        return res;
    }
private:
    const long long K = 10000;  // 偏移系数
    int res = 0;                // 结果值

    /**
     * 获取以0节点为根时的正确数
     */
    int dfs_GetZeroRootCorrect(int node, int fa, vector<vector<int>> &neighbors, unordered_set<long long> &guessesIds)
    {
        int correct = 0;  // 统计正确数
        for (auto &child : neighbors[node]) {
            if (child != fa) {
                correct += (guessesIds.count(node * K + child));  // 当前父子节点对(node,child)存在于guessIds中，猜中
                correct += dfs_GetZeroRootCorrect(child, node, neighbors, guessesIds);  // 递归搜索子树情况
            }
        }
        return correct;
    }

    /**
     * 通过换根策略，得到每个节点为根时的正确数
     */
    void dfs_ChangeRoot(
        int node, int fa, int k, int correct, vector<vector<int>> &neighbors, unordered_set<long long> &guessesIds)
    {
        if (correct >= k)
            res++;  // 以node为根节点的猜中数correct大于等于k，是一种可能的情况
        for (auto &child : neighbors[node]) {
            if (child != fa) {
                // 递归搜索子节点
                // 如果当前(node,child)存在，当前是猜中的，下一步递归到子节点，子节点变成了根节点，就是猜错了，因此正确数要-1
                // 如果(child,node)存在，下一步递归到子节点，子节点变成了根节点，原本父节点就变成子节点，就是猜对了，因此正确数要+1
                dfs_ChangeRoot(child,
                    node,
                    k,
                    correct - guessesIds.count(node * K + child) + guessesIds.count(child * K + node),
                    neighbors,
                    guessesIds);
            }
        }
    }

};
int main()
{
    // edges = [[0,1],[1,2],[1,3],[4,2]], guesses = [[1,3],[0,1],[1,0],[2,4]], k = 3 输出：3
    vector<vector<int>> edges0 = {{0, 1}, {1, 2}, {1, 3}, {4, 2}};
    vector<vector<int>> guesses0 = {{1, 3}, {0, 1}, {1, 0}, {2, 4}};
    int k0 = 3;
    int res0 = Solution().rootCount(edges0, guesses0, k0);
    cout << "res0: " << res0 << endl;

    // edges = [[0,1],[1,2],[2,3],[3,4]], guesses = [[1,0],[3,4],[2,1],[3,2]], k = 1 输出：5
    vector<vector<int>> edges1 = {{0, 1}, {1, 2}, {2, 3}, {3, 4}};
    vector<vector<int>> guesses1 = {{1, 0}, {3, 4}, {2, 1}, {3, 2}};
    int k1 = 1;
    int res1 = Solution().rootCount(edges1, guesses1, k1);
    cout << "res1: " << res1 << endl;
}