DFS_medium_5406_收集树上所有苹果的最少时间

题目描述

思路

方法一：DFS

def minTime(n: int, edges: List[List[int]], hasApple: List[bool]) -> int:
  data_dict = {}
  for i in range(n):
    data_dict[i] = []
  for one in edges:
    data_dict[one[0]].append(one[1])
  
  # print(data_dict)
  def dfs(i):
    key = hasApple[i]
    for index in data_dict[i]:
      key |= dfs(index)
    if not key:
      del data_dict[i]
    return key

  dfs(0)
  # print(data_dict)
  return max(0, 2*(len(data_dict) - 1))

class Solution:
    def minTime(self, n: int, edges: List[List[int]], hasApple: List[bool]) -> int:
        data_dict = {}
        for one in edges:
            if one[0] in data_dict:
                data_dict[one[0]].append(one[1])
            else:
                data_dict[one[0]] = [one[1]]
        # dfs
        def dfs(i):
            cnt = 0
            if i in data_dict:
                for child in data_dict[i]:
                    tmp = dfs(child)
                    cnt += tmp
                    # 如果此孩子节点为True（苹果）或者 该节点的孩子节点统计的数量不为0（有苹果）
                    if hasApple[child] or tmp:  
                        cnt += 2
            return cnt
        return dfs(0)

方法二：类似于Union-Find方法

• 层层向上统计父节点，并用 hash set 存储，用来去重。

class Solution:
    def minTime(self, n: int, edges: List[List[int]], hasApple: List[bool]) -> int:
        data_dict = {}
        for one in edges:
          data_dict[one[1]] = one[0]
        
        print(data_dict)
        cnt = 0
        father_node = set()
        for index, ele in enumerate(hasApple):
          if ele:
            while index not in father_node and index in data_dict:
              cnt += 2
              father_node.add(index)
              index = data_dict[index]
        return cnt