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

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 搜索即可~

---