863_二叉树中所有距离为_K_的结点

类型：树

• 863. 二叉树中所有距离为 K 的结点 💛

  https://leetcode-cn.com/problems/all-nodes-distance-k-in-binary-tree/

  ❓ 找到二叉树中距目标结点 target 距离为 K 的所有结点。返回结点值列表。

  💡 DFS

  如果 target 节点在 root 节点的左子树中，且 target 节点深度为 3，那所有 root 节点右子树中深度为 K - 3 的节点到 target 的距离就都是 K。

  那么我们定义深度优先搜索方法 dfs(node)，该方法会返回 node 到 target 的距离。在 dfs(node) 中处理以下四种情况：

  • node == target，则把子树中所有距 node 距离为 k 的结点加入答案。
  • target 在 node 左子树中，假设 target 与 node 距离为 L+1，则找出 node 右子树中所有距离为 K - L - 1 的结点加入答案。
  • target 在 node 右子树中，与在左子树中的处理方法一样。
  • target 不在 node 子树中，不用处理。

  class Solution(object):
      def distanceK(self, root, target, K):
          ans = []
  
          # Return distance from node to target if exists, else -1
          # Vertex distance: the # of vertices on the path from node to target
          def dfs(node):
              if not node:
                  return -1
              elif node is target:
                  subtree_add(node, 0)
                  return 1
              else:
                  L, R = dfs(node.left), dfs(node.right)
                  if L != -1:
                      if L == K: ans.append(node.val)
                      subtree_add(node.right, L + 1)
                      return L + 1
                  elif R != -1:
                      if R == K: ans.append(node.val)
                      subtree_add(node.left, R + 1)
                      return R + 1
                  else:
                      return -1
  
          # Add all nodes 'K - dist' from the node to answer.
          def subtree_add(node, dist):
              if not node:
                  return
              elif dist == K:
                  ans.append(node.val)
              else:
                  subtree_add(node.left, dist + 1)
                  subtree_add(node.right, dist + 1)
  
          dfs(root)
          return ans

  时间复杂度：O(n)，空间复杂度：O(n)