863. All Nodes Distance K in Binary Tree

题目描述

leetcode-cn.com/problems/al…

分析

找到距离 target 为 k 的所有节点

从 target 开始遍历，走 k 步

解题思路

记录下所有节点的父节点，从 target 开始遍历，需要注意：

遍历的路径分为三部分：向上，遍历 left，遍历 right

过程

统计所有 parent node

const findParents = (node) => {
    if (node.left != null) {
      parents.set(node.left.val, node)
      findParents(node.left)
    }
    if (node.right != null) {
      parents.set(node.right.val, node)
      findParents(node.right)
    }
  }

用一个 map 进行统计，从根节点开始进行 dfs

寻找所有深度为 k 的结点

const findAns = (node, from, depth, k) => {
    if (node == null) {
      return
    }
    if (depth === k) {
      ans.push(node.val)
      return
    }
    if (node.left !== from) {
      findAns(node.left, node, depth + 1, k)
    }
    if (node.right !== from) {
      findAns(node.right, node, depth + 1, k)
    }
    if (parents.get(node.val) !== from) {
      findAns(parents.get(node.val), node, depth + 1, k)
    }
  }

从 target 开始通过递归开始遍历，需要注意的是：

需要有一个 from 参数记录遍历的上一个节点，这是为了不重复遍历

代码

var distanceK = function (root, target, k) {
const parents = new Map()
const ans = []

const findParents = (node) => {
  if (node.left != null) {
    parents.set(node.left.val, node)
    findParents(node.left)
  }
  if (node.right != null) {
    parents.set(node.right.val, node)
    findParents(node.right)
  }
}

// 从 root 出发 DFS，记录每个结点的父结点
findParents(root)

const findAns = (node, from, depth, k) => {
  if (node == null) {
    return
  }
  if (depth === k) {
    ans.push(node.val)
    return
  }
  if (node.left !== from) {
    findAns(node.left, node, depth + 1, k)
  }
  if (node.right !== from) {
    findAns(node.right, node, depth + 1, k)
  }
  if (parents.get(node.val) !== from) {
    findAns(parents.get(node.val), node, depth + 1, k)
  }
}
// 从 target 出发 DFS，寻找所有深度为 k 的结点
findAns(target, null, 0, k)

return ans
}