lc2581.统计可能的树根数目

题目链接：

题解

如果已知 fa 为根时，满足猜测的答案是 now，那么如果 nx 为根，fa 为叶子，对答案的贡献就是 (fa, nx) 和 (nx, fa)。

第一次以 0 为根，计算满足的答案，之后再 dfs，每次值影响两个节点的关系。

参考代码

#
# @lc app=leetcode.cn id=2581 lang=python3
#
# [2581] 统计可能的树根数目
#

# @lc code=start
class Solution:
    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:
        g = [[] for i in range(len(edges) + 1)]
        for a in edges:
            g[a[0]].append(a[1])
            g[a[1]].append(a[0])
        guesse = {(a[0], a[1]) for a in guesses}
        ans = 0
        base = 0
        def dfs(rt, fa):

            if (fa, rt) in guesse:
                nonlocal base
                base += 1
            for nx in g[rt]:
                if nx == fa:
                    continue
                dfs(nx, rt)

        dfs(0, -1)
        # print(base)

        def check(rt, fa, now):
            nonlocal ans
            if (now >= k):
                ans += 1
            for nx in g[rt]:
                if nx == fa:
                    continue
                tmp = now
                if (rt, nx) in guesse:
                    tmp -= 1
                if (nx, rt) in guesse:
                    tmp += 1
                check(nx, rt, tmp)

        check(0, -1, base)
        return ans
# @lc code=end