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

原创

haocold 2022-10-24 19:20:23 博主文章分类：算法 ©著作权

文章标签 swift 算法二叉树深度优先结点 文章分类 虚拟化云计算

©著作权归作者所有：来自51CTO博客作者haocold的原创作品，请联系作者获取转载授权，否则将追究法律责任

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) -> [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
}