# 题目地址(1261. 在受污染的二叉树中查找元素) https://leetcode-cn.com/problems/find-elements-in-a-contaminated-binary-tree/submissions/ ## 题目描述 ``` 给出一个满足下述规则的二叉树: 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 是否存在于还原后的二叉树中并返回结果。   示例 1: ![](https://tva1.sinaimg.cn/large/006tNbRwgy1gasy4qroxoj308w03b3yi.jpg) 输入: ["FindElements","find","find"] [[[-1,null,-1]],[1],[2]] 输出: [null,false,true] 解释: FindElements findElements = new FindElements([-1,null,-1]); findElements.find(1); // return False findElements.find(2); // return True 示例 2: ![](https://tva1.sinaimg.cn/large/006tNbRwgy1gasy5mlo3mj30b405iwep.jpg) 输入: ["FindElements","find","find","find"] [[[-1,-1,-1,-1,-1]],[1],[3],[5]] 输出: [null,true,true,false] 解释: FindElements findElements = new FindElements([-1,-1,-1,-1,-1]); findElements.find(1); // return True findElements.find(3); // return True findElements.find(5); // return False 示例 3: ![](https://tva1.sinaimg.cn/large/006tNbRwgy1gasy5sr25yj308i07maa8.jpg) 输入: ["FindElements","find","find","find","find"] [[[-1,null,-1,-1,null,-1]],[2],[3],[4],[5]] 输出: [null,true,false,false,true] 解释: FindElements findElements = new FindElements([-1,null,-1,-1,null,-1]); findElements.find(2); // return True findElements.find(3); // return False findElements.find(4); // return False findElements.find(5); // return True   提示: TreeNode.val == -1 二叉树的高度不超过 20 节点的总数在 [1, 10^4] 之间 调用 find() 的总次数在 [1, 10^4] 之间 0 <= target <= 10^6 ``` ## 暴力法 ### 思路 最简单想法就是递归建立树,然后 find 的时候递归查找即可,代码也很简单。 ### 代码 Pythpn Code: ```python # Definition for a binary tree node. # class TreeNode: # def __init__(self, x): # self.val = x # self.left = None # self.right = None class FindElements: node = None def __init__(self, root: TreeNode): def recover(node): if not node: return node; if node.left: node.left.val = 2 * node.val + 1 if node.right: node.right.val = 2 * node.val + 2 recover(node.left) recover(node.right) return node root.val = 0 self.node = recover(root) def find(self, target: int) -> bool: def findInTree(node, target): if not node: return False if node.val == target: return True return findInTree(node.left, target) or findInTree(node.right, target) return findInTree(self.node, target) # Your FindElements object will be instantiated and called as such: # obj = FindElements(root) # param_1 = obj.find(target) ``` 上述代码会超时,我们来考虑优化。 ## 空间换时间 ### 思路 上述代码会超时,我们考虑使用空间换时间。 建立树的时候,我们将所有值存到一个集合中去。当需要 find 的时候,我们直接查找 set 即可,时间复杂度 O(1)。 ### 代码 ```python # Definition for a binary tree node. # class TreeNode: # def __init__(self, x): # self.val = x # self.left = None # self.right = None class FindElements: def __init__(self, root: TreeNode): # set 不能放在init外侧。 因为测试用例之间不会销毁FindElements的变量 self.seen = set() def recover(node): if not node: return node; if node.left: node.left.val = 2 * node.val + 1 self.seen.add(node.left.val) if node.right: node.right.val = 2 * node.val + 2 self.seen.add(node.right.val) recover(node.left) recover(node.right) return node root.val = 0 self.seen.add(0) self.node = recover(root) def find(self, target: int) -> bool: return target in self.seen # Your FindElements object will be instantiated and called as such: # obj = FindElements(root) # param_1 = obj.find(target) ``` 这种解法可以 AC,但是在数据量非常大的时候,可能 MLE,我们继续考虑优化。 ## 二进制法 ### 思路 这是一种非常巧妙的做法。 如果我们把树中的数全部加 1 会怎么样? ![](https://tva1.sinaimg.cn/large/006tNbRwly1gasypfuvuvj30rs0kudjr.jpg) (图参考 https://leetcode.com/problems/find-elements-in-a-contaminated-binary-tree/discuss/431229/Python-Special-Way-for-find()-without-HashSet-O(1)-Space-O(logn)-Time) 仔细观察发现,每一行的左右子树分别有不同的前缀: ![](https://tva1.sinaimg.cn/large/006tNbRwgy1gasz0x09koj312y0sgnnt.jpg) Ok,那么算法就来了。为了便于理解,我们来举个具体的例子,比如 target 是 9,我们首先将其加 1,二进制表示就是 1010。不考虑第一位,就是 010,我们只要: - 0 向左 👈 - 1 向右 👉 - - 0 向左 👈 就可以找到 9 了。 > 0 表示向左 , 1 表示向右 ### 代码 ```python # Definition for a binary tree node. # class TreeNode: # def __init__(self, x): # self.val = x # self.left = None # self.right = None class FindElements: node = None def __init__(self, root: TreeNode): def recover(node): if not node: return node; if node.left: node.left.val = 2 * node.val + 1 if node.right: node.right.val = 2 * node.val + 2 recover(node.left) recover(node.right) return node root.val = 0 self.node = recover(root) def find(self, target: int) -> bool: node = self.node for bit in bin(target+1)[3:]: node = node and (node.left, node.right)[int(bit)] return bool(node) # Your FindElements object will be instantiated and called as such: # obj = FindElements(root) # param_1 = obj.find(target) ``` ## 关键点解析 - 空间换时间 - 二进制思维 - 将 target + 1