-
Notifications
You must be signed in to change notification settings - Fork 0
/
AmountofTimeforBinaryTreetoBeInfected.py
44 lines (35 loc) · 1.19 KB
/
AmountofTimeforBinaryTreetoBeInfected.py
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
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, val=0, left=None, right=None):
# self.val = val
# self.left = left
# self.right = right
class Solution:
def amountOfTime(self, root: Optional[TreeNode], start: int) -> int:
# edge 2 up
# graph
# unique number = no cycle
def dfs(graph, visited, start):
visited[start] = True
ans = 0
for n in graph[start]:
if(not visited[n]):
ans = max(ans, dfs(graph, visited, n)+1)
return ans
N = 100001
visited = [False]*N
graph = [[] for _ in range(N)]
# make graph
stk = [root]
while(len(stk)):
node = stk.pop()
if(node.left):
graph[node.val].append(node.left.val)
graph[node.left.val].append(node.val)
stk.append(node.left)
if(node.right):
graph[node.val].append(node.right.val)
graph[node.right.val].append(node.val)
stk.append(node.right)
# dfs
return dfs(graph, visited, start)