[Leetcode] 二叉树中所有距离为 K 的结点

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

思路

1. 根据题意,我们首先需要找到target所在的节点, 并以其为原点, 向左右字节点, 父节点 进行查找. 故我们需要考虑建立一个父子节点的关系标记方便后续的查找. 故我们通过引入map 和 dfs遍历来构建子节点和父节点的的关系(以 val为key, 因为题目中说明,树中的每个节点之值唯一)
2. 当我们找到target时, 以target节点为根结点, 重新进行一次dfs遍历, 当遍历的深度 depth === k 时, 即为我们需要找到的节点. (因为只需要寻找距离为k, 当找到节点时,即可停止该条线路的遍历

代码

/**
 * Definition for a binary tree node.
 */
function TreeNode(val) {
  this.val = val
  this.left = this.right = null
}

/**
 * 节点值唯一. 故我们可以以节点值标记每一个节点
 * 我们首先需要找到目标 target. 以目标target开始往上或往下寻找距离点
 * 我们可以通过记录 节点之间的父子关系, 实现向上寻找
 * @param {TreeNode} root
 * @param {TreeNode} target
 * @param {number} k
 * @return {number[]}
 */
var distanceK = function (root, target, k) {
  const map = new Map()
  const dfs = (node) => {
    if (node.left !== null) {
      map.set(node.left.val, node)
      dfs(node.left)
    }
    if (node.right !== null) {
      map.set(node.right.val, node)
      dfs(node.right)
    }
  }
  dfs(root)

  const result = []
  /**
   * @param {*} node 当前查询的节点
   * @param {*} from 是从哪个节点查询过来的, 避免重复查询与陷入死循环
   * @param {*} depth 遍历的深度 记录与target的距离
   */
  const find = (node, from, depth) => {
    if (!node) return
    if (depth === k) {
      result.push(node.val)
      return
      // 提前返回结果, 减少无必要的遍历
    }
    if (node.left !== from) find(node.left, node, depth + 1)
    if (node.right !== from) find(node.right, node, depth + 1)

    const parent = map.get(node.val)
    if (parent !== from) find(parent, node, depth + 1)
  }
  find(target, null, 0)
  return result
}