2581.统计可能的树根数目

链接：2581.统计可能的树根数目
难度：Hard
标签：树、深度优先搜索、数组、哈希表、动态规划
简介：给你二维整数数组 edges ，Bob 的所有猜测和整数 k ，请你返回可能成为树根的 节点数目 。如果没有这样的树，则返回 0。

题解 1 - python

• 编辑时间：2024-02-29
• 执行用时：425ms
• 内存消耗：97.89MB
• 编程语言：python
• 解法介绍：先统计以0为根的猜对个数， 再对每个节点为根进行dfs。

class Solution:
    def rootCount(self, edges: List[List[int]], guesses: List[List[int]], k: int) -> int:
        nodes = [[] for _ in range(len(edges) + 1)]
        for n1, n2 in edges:
            nodes[n1].append(n2)
            nodes[n2].append(n1)
        s = {(n1, n2) for n1, n2 in guesses}
        def dfs(node: int, parent: int) -> int:
            ans = 0
            for child in nodes[node]:
                if child != parent:
                    ans += (node, child) in s
                    ans += dfs(child, node)
            return ans
        def reroot(node: int, parent: int, cnt: int) -> int:
            ans = cnt >= k
            for child in nodes[node]:
                if child != parent:
                    ans += reroot(child, node, cnt + ((child, node) in s) - ((node, child) in s))
            return ans
        return reroot(0, -1, dfs(0, -1))