2385. 感染二叉树需要的总时间

感染二叉树需要的总时间

• 难度: medium
• 原始链接: https://leetcode.cn/problems/amount-of-time-for-binary-tree-to-be-infected
• 标签: 广度优先遍历, 深度优先遍历

解法一: 广度优先遍历, 深度优先遍历

go

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func amountOfTime(root *TreeNode, start int) int {
  relations := make(map[int][]int)
  var searchParent func(cur, pre *TreeNode)
  searchParent = func(cur, pre *TreeNode) {
    list := []int{}
    if pre != nil {
      list = append(list, pre.Val)
    }
    if cur.Left != nil {
      list = append(list, cur.Left.Val)
      searchParent(cur.Left, cur)
    }
    if cur.Right != nil {
      list = append(list, cur.Right.Val)
      searchParent(cur.Right, cur)
    }
    relations[cur.Val] = list
  }
  searchParent(root, nil)
  queue := []int{start}
  visit := map[int]bool{start: true}
  ans := 0
  for {
    size := len(queue)
    for i := 0; i < size; i++ {
      for _, nodeVal := range relations[queue[i]] {
        if !visit[nodeVal] {
          visit[nodeVal] = true
          queue = append(queue, nodeVal)
        }
      }
    }
    queue = queue[size:]
    if len(queue) == 0 {
      break
    }
    ans++
  }
  return ans
}