aboutsummaryrefslogtreecommitdiff
path: root/amount_time_bt_infect.py
blob: 9592126888a6ca5a603b95f3c017a5abfe5871a9 (plain)
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
49
50
51
52
53
54
55
56
57
58
59
# 2385. Amount of Time for Binary Tree to Be Infected

"""
given the 'root' of a binary tree with unique values, and an integer 'start'.
at minute 0, an infection starts from the node with value 'start' each
minute, a node becomes infected if the node is currently uninfected, or the
node is adjacent to an infected node. return the number of minutes needed for
the entire tree to be infected.
"""


# Definition for a binary tree node.
class TreeNode(object):
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution(object):
    def amountOfTime(self, root, start):
        """
        :type root: Optional[TreeNode]
        :type start: int
        :rtype: int
        """

        def dfs(node):
            if node is None:
                return
            if node.left:
                graph[node.val].append(node.left.val)
                graph[node.left.val].append(node.val)
            if node.right:
                graph[node.val].append(node.right.val)
                graph[node.right.val].append(node.val)
            dfs(node.left)
            dfs(node.right)

        graph = defaultdict(list)
        dfs(root)
        visited = set()
        queue = deque([start])
        time = -1
        while queue:
            time += 1
            for _ in range(len(queue)):
                current_node = queue.popleft()
                visited.add(current_node)
                for neighbor in graph[current_node]:
                    if neighbor not in visited:
                        queue.append(neighbor)
        return time


if __name__ == "__main__":
    obj = Solution()
    print(obj.amountOfTime(root=[1, 5, 3, null, 4, 10, 6, 9, 2], start=3))
    print(obj.amountOfTime(root=[1], start=1))