LC 2581. 统计可能的树根数目

https://leetcode.cn/problems/count-number-of-possible-root-nodes/

题解解释得非常清楚了：

使用"假设0为根"来计算基准值
如果之前假设x为根，现在假设y为根的话，那么
- 删除(x,y)
- 增加(y,x)
不过前提是x,y之间有联通边

"换根DP", 这个解法之前还没有听说过。

class Solution:
    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:
        n = len(edges) + 1
        adj = [[] for _ in range(n)]
        for x, y in edges:
            adj[x].append(y)
            adj[y].append(x)

        ss = {(x, y) for (x, y) in guesses}

        def dfs(x, p):
            r = 0
            for y in adj[x]:
                if y == p: continue
                r += ((x, y) in ss)
                r += dfs(y, x)
            return r

        base = dfs(0, -1)

        def reroot(x, p, now):
            r = 0
            if now >= k: r += 1
            for y in adj[x]:
                if y == p: continue
                r += reroot(y, x, (now - ((x, y) in ss) + ((y, x) in ss)))
            return r

        ans = reroot(0, -1, base)
        return ans