1443. Minimum Time to Collect All Apples in a Tree

My first thought was to use dfs to iterate the tree 2 times: first to check if node is worth visiting then to collect apple

The below implementation achieve the 2 goals in one iteration: if childTime > 0 or hasApple[child]: is used to denote whether current node and it's children contribute to the collection time

pass in a parent node to avoid revisiting parent
a child is worth visiting if itself has apple or there is apple in the tree where root is the child

class Solution:
    def minTime(self, n: int, edges: List[List[int]], hasApple: List[bool]) -> int:
        adj = [[] for _ in range(n)]
        for i, j in edges:
            adj[i].append(j)
            adj[j].append(i)

        return self.getTime(adj, hasApple, -1, 0)
    
    def getTime(self, adj, hasApple, parent, cur): # time to collect apple in a tree where cur is root
        total = 0
        for child in adj[cur]:
            # return 0 for leave node
            if child == parent:
                continue
            childTime = self.getTime(adj, hasApple, cur, child)
            # case 1: child has apple, there might/might not be apple in the tree where child is root (childTime =/> 0)
            # case 2: child has no apple, but there are apples in the tree where child is the root
            if childTime > 0 or hasApple[child]:
                total += childTime + 2
        return total

1443. Minimum Time to Collect All Apples in a Tree - cocoder39/coco39_LC GitHub Wiki