1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48
| # # @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
|