Leetcode 1554 Solution

This article provides solution to leetcode question 1554 (minimum-time-to-collect-all-apples-in-a-tree)

https://leetcode.com/problems/minimum-time-to-collect-all-apples-in-a-tree

Solution

class Solution:
    def minTime(self, n: int, edges: List[List[int]], hasApple: List[bool]) -> int:
        graph = collections.defaultdict(list)
        for src, dst in edges:
            graph[src].append(dst)
            graph[dst].append(src)

        apple_node_set = {i for i, apple in enumerate(hasApple) if apple}

        visited = set()

        def dfs(node):
            if node in visited:
                return False, 0

            visited.add(node)

            subtree_has_apple = node in apple_node_set
            min_path = 0

            for child in graph[node]:
                child_subtree_has_apple, child_min_path = dfs(child)
                subtree_has_apple |= child_subtree_has_apple
                min_path += child_min_path + int(child_subtree_has_apple)

            return subtree_has_apple, min_path

        return dfs(0)[1] * 2