paradox64ce

lc

Sep 4th, 2021
611
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. class Solution:
  2.     def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
  3.         ans = [0] * n
  4.         adj = [[] for x in range(n)] # why does [[]] * n not work ??
  5.         sub = [0] * n
  6.         depth = 0
  7.        
  8.         for u, v in edges:
  9.             # print(u, v)
  10.             adj[u].append(v)
  11.             adj[v].append(u)
  12.        
  13.         # for v in adj[0]:
  14.         #     print(v)
  15.        
  16.         def dfs1(u, p, depth):
  17.             # global depth
  18.             cnt = 1
  19.             ans[0] += depth
  20.             for v in adj[u]:
  21.                 if v!=p:
  22.                     depth += 1
  23.                     cnt += dfs1(v, u, depth)
  24.                     depth -= 1
  25.             sub[u] = cnt
  26.             return cnt
  27.        
  28.         def dfs2(u, p):
  29.             for v in adj[u]:
  30.                 if v!=p:
  31.                     ans[v] = ans[u] + n - 2*sub[v]
  32.                     dfs2(v, u)
  33.        
  34.         dfs1(0,-1, 0)
  35.         dfs2(0,-1)
  36.        
  37.         return ans
  38.        
  39.            
RAW Paste Data