【二叉树】二叉树中所有距离为 K 的结点

本文介绍: 题目给定一个二叉树（具有根结点 root）一个目标结点 target，和一个整数值 k 返回到目标结点 target 距离为 k 的所有结点的值的列表答案可以以任何顺序返回

0x00 题目

给定一个二叉树（具有根结点 root）
一个目标结点 target，和一个整数值 k
返回到目标结点 target 距离为 k 的
所有结点的值的列表
 答案可以以任何顺序返回

0x01 思路

如果 target 是 根节点
通过深度优先方法遍历
 找出深度为 k 的节点即可

当 target 不是 根节点 时
如何根据当前节点找出父节点呢？
可以通过字典来解决
以当前节点值为 key
父节点为 value

所以要先遍历整棵树
然后再查找

0x02 解法

语言：Swift

树节点：TreeNode

public class TreeNode {
    public var val: Int
    public var left: TreeNode?
    public var right: TreeNode?
    public init() { self.val = 0; self.left = nil; self.right = nil; }
    public init(_ val: Int) { self.val = val; self.left = nil; self.right = nil; }
    public init(_ val: Int, _ left: TreeNode?, _ right: TreeNode?) {
        self.val = val
        self.left = left
        self.right = right
    }
}

解法：

func distanceK(_ root: TreeNode?, _ target: TreeNode?, _ k: Int) -&gt; [Int] {
    var res: [Int] = []
    var parent: [Int:TreeNode] = [:]
    
    // 查找出所有节点的父节点
    func findParent(_ root: TreeNode?) {
        guard let r = root else { return }
        if r.left != nil {
            parent[r.left!.val] = r;
            findParent(r.left)
        }
        if r.right != nil {
            parent[r.right!.val] = r;
            findParent(r.right)
        }
    }
    
    // 查找
    func find(_ root: TreeNode?, _ pre: TreeNode?, _ dep: Int, _ k: Int) {
        guard let r = root else { return }
        // 深度相同，找到
        if dep == k {
            res.append(r.val)
            return
        }
        
        // 往左找
        if r.left !== pre {
            find(r.left, r, dep+1, k)
        }
        
        // 往右找
        if r.right !== pre {
            find(r.right, r, dep+1, k)
        }
        
        // 往父节点找，过滤目标节点的子节点
        if parent[r.val] !== pre {
            find(parent[r.val], r, dep+1, k)
        }
    }
    
    findParent(root)
    // 从目标节点开始找
    find(target, nil, 0, k)
    
    return res
}