【二叉树】在受污染的二叉树中查找元素

0x00 题目

给出一个满足下述规则的二叉树：
​​​root.val == 0​​​ 如果 ​​treeNode.val == x​​ 且 ​​treeNode.left != null​​ 那么 ​​treeNode.left.val == 2 * x + 1​​ 如果 ​​treeNode.val == x​​ 且 ​​treeNode.right != null​​ 那么 ​​treeNode.right.val == 2 * x + 2​​ 现在这个二叉树受到「污染」，所有的 ​​treeNode.val​​ 都变成了 ​​-1​​

请你先​​还原​​二叉树，然后实现 FindElements 类：

​​FindElements(TreeNode* root)​​​ 用受污染的二叉树初始化对象，你需要先把它还原
​​​bool find(int target)​​​ 判断目标值 ​​target​​ 是否存在于还原后的二叉树中并返回结果

0x01 思路

还原二叉树，根据​​子​​​节点与​​根​​​节点的关系
即可​​​计算​​​出来
是否包含，使用​​​集合​​​来判断即可
在恢复二叉树的过程中
把值添加到集合即可

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

解法：

class FindElements {
    var set = Set<Int>()

    init(_ root: TreeNode?) {
        root?.val = 0
        set.insert(0)
        
        // 修复二叉树
        fix(root)
    }

    func find(_ target: Int) -> Bool {
        return set.contains(target)
    }

    private func fix(_ root: TreeNode?) {
        // 左节点不为空
        if let left = root?.left {
            let val = (root!.val << 1) + 1
            set.insert(val)
            left.val = val
            fix(left)
        }
        // 右节点不为空        
        if let right = root?.right {
            let val = (root!.val << 1) + 2
            set.insert(val)
            right.val = val
            fix(right)
        }
    }
}

0x03 我的作品

欢迎体验我的作品之一：​​小五笔 86 版​​​ 五笔学习好帮手！
​​App Store​​ 搜索即可~