Question

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

English Version

题目描述

给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 k 。

返回到目标结点 target 距离为 k 的所有结点的值的列表。 答案可以以 任何顺序 返回。

 

示例 1：

输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
输出：[7,4,1]
解释：所求结点为与目标结点（值为 5）距离为 2 的结点，值分别为 7，4，以及 1

示例 2:

输入: root = [1], target = 1, k = 3
输出: []

 

提示:

• 节点数在 [1, 500] 范围内
• 0 <= Node.val <= 500
• Node.val 中所有值 不同
• 目标结点 target 是树上的结点。
• 0 <= k <= 1000

 

解法

方法一：DFS + 哈希表

我们先用 DFS 遍历整棵树，记录每个结点的父结点，然后从目标结点开始，向上、向下分别搜索距离为 $k$ 的结点，添加到答案数组中。

时间复杂度 $O(n)$，空间复杂度 $O(n)$。其中 $n$ 是二叉树的结点数。

# Definition for a binary tree node.
# class TreeNode:
#   def __init__(self, x):
#     self.val = x
#     self.left = None
#     self.right = None


class Solution:
  def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
    def parents(root, prev):
      nonlocal p
      if root is None:
        return
      p[root] = prev
      parents(root.left, root)
      parents(root.right, root)

    def dfs(root, k):
      nonlocal ans, vis
      if root is None or root.val in vis:
        return
      vis.add(root.val)
      if k == 0:
        ans.append(root.val)
        return
      dfs(root.left, k - 1)
      dfs(root.right, k - 1)
      dfs(p[root], k - 1)

    p = {}
    parents(root, None)
    ans = []
    vis = set()
    dfs(target, k)
    return ans

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *   int val;
 *   TreeNode left;
 *   TreeNode right;
 *   TreeNode(int x) { val = x; }
 * }
 */
class Solution {
  private Map<TreeNode, TreeNode> p;
  private Set<Integer> vis;
  private List<Integer> ans;

  public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
    p = new HashMap<>();
    vis = new HashSet<>();
    ans = new ArrayList<>();
    parents(root, null);
    dfs(target, k);
    return ans;
  }

  private void parents(TreeNode root, TreeNode prev) {
    if (root == null) {
      return;
    }
    p.put(root, prev);
    parents(root.left, root);
    parents(root.right, root);
  }

  private void dfs(TreeNode root, int k) {
    if (root == null || vis.contains(root.val)) {
      return;
    }
    vis.add(root.val);
    if (k == 0) {
      ans.add(root.val);
      return;
    }
    dfs(root.left, k - 1);
    dfs(root.right, k - 1);
    dfs(p.get(root), k - 1);
  }
}

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *   int val;
 *   TreeNode *left;
 *   TreeNode *right;
 *   TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
  unordered_map<TreeNode*, TreeNode*> p;
  unordered_set<int> vis;
  vector<int> ans;

  vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
    parents(root, nullptr);
    dfs(target, k);
    return ans;
  }

  void parents(TreeNode* root, TreeNode* prev) {
    if (!root) return;
    p[root] = prev;
    parents(root->left, root);
    parents(root->right, root);
  }

  void dfs(TreeNode* root, int k) {
    if (!root || vis.count(root->val)) return;
    vis.insert(root->val);
    if (k == 0) {
      ans.push_back(root->val);
      return;
    }
    dfs(root->left, k - 1);
    dfs(root->right, k - 1);
    dfs(p[root], k - 1);
  }
};

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *   Val int
 *   Left *TreeNode
 *   Right *TreeNode
 * }
 */
func distanceK(root *TreeNode, target *TreeNode, k int) []int {
  p := make(map[*TreeNode]*TreeNode)
  vis := make(map[int]bool)
  var ans []int
  var parents func(root, prev *TreeNode)
  parents = func(root, prev *TreeNode) {
    if root == nil {
      return
    }
    p[root] = prev
    parents(root.Left, root)
    parents(root.Right, root)
  }
  parents(root, nil)
  var dfs func(root *TreeNode, k int)
  dfs = func(root *TreeNode, k int) {
    if root == nil || vis[root.Val] {
      return
    }
    vis[root.Val] = true
    if k == 0 {
      ans = append(ans, root.Val)
      return
    }
    dfs(root.Left, k-1)
    dfs(root.Right, k-1)
    dfs(p[root], k-1)
  }
  dfs(target, k)
  return ans
}

方法二

# Definition for a binary tree node.
# class TreeNode:
#   def __init__(self, x):
#     self.val = x
#     self.left = None
#     self.right = None


class Solution:
  def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
    def dfs1(root, fa):
      if root is None:
        return
      p[root] = fa
      dfs1(root.left, root)
      dfs1(root.right, root)

    def dfs2(root, fa, k):
      if root is None:
        return
      if k == 0:
        ans.append(root.val)
        return
      for nxt in (root.left, root.right, p[root]):
        if nxt != fa:
          dfs2(nxt, root, k - 1)

    p = {}
    dfs1(root, None)
    ans = []
    dfs2(target, None, k)
    return ans