# 1261. 在受污染的二叉樹中查找元素 # 題目地址（1261. 在受污染的二叉樹中查找元素） <https://leetcode-cn.com/problems/find-elements-in-a-contaminated-binary-tree/> ## 題目描述 ``` <pre class="calibre18">``` 給出一個滿足下述規則的二叉樹： 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：  輸入： ["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：  輸入： ["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：  輸入： ["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: ``` <pre class="calibre18">``` <span class="hljs-title"># Definition for a binary tree node.</span> <span class="hljs-title"># class TreeNode:</span> <span class="hljs-title"># def __init__(self, x):</span> <span class="hljs-title"># self.val = x</span> <span class="hljs-title"># self.left = None</span> <span class="hljs-title"># self.right = None</span> <span class="hljs-class"><span class="hljs-keyword">class</span> <span class="hljs-title">FindElements</span>:</span> node = <span class="hljs-keyword">None</span> <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">__init__</span><span class="hljs-params">(self, root: TreeNode)</span>:</span> <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">recover</span><span class="hljs-params">(node)</span>:</span> <span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> node: <span class="hljs-keyword">return</span> node; <span class="hljs-keyword">if</span> node.left: node.left.val = <span class="hljs-params">2</span> * node.val + <span class="hljs-params">1</span> <span class="hljs-keyword">if</span> node.right: node.right.val = <span class="hljs-params">2</span> * node.val + <span class="hljs-params">2</span> recover(node.left) recover(node.right) <span class="hljs-keyword">return</span> node root.val = <span class="hljs-params">0</span> self.node = recover(root) <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">find</span><span class="hljs-params">(self, target: int)</span> -> bool:</span> <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">findInTree</span><span class="hljs-params">(node, target)</span>:</span> <span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> node: <span class="hljs-keyword">return</span> <span class="hljs-keyword">False</span> <span class="hljs-keyword">if</span> node.val == target: <span class="hljs-keyword">return</span> <span class="hljs-keyword">True</span> <span class="hljs-keyword">return</span> findInTree(node.left, target) <span class="hljs-keyword">or</span> findInTree(node.right, target) <span class="hljs-keyword">return</span> findInTree(self.node, target) <span class="hljs-title"># Your FindElements object will be instantiated and called as such:</span> <span class="hljs-title"># obj = FindElements(root)</span> <span class="hljs-title"># param_1 = obj.find(target)</span> ``` ``` 上述代碼會超時，我們來考慮優化。 ## 空間換時間 ### 思路 上述代碼會超時，我們考慮使用空間換時間。 建立樹的時候，我們將所有值存到一個集合中去。當需要 find 的時候，我們直接查找 set 即可，時間復雜度 O(1)。 ### 代碼 ``` <pre class="calibre18">``` <span class="hljs-title"># Definition for a binary tree node.</span> <span class="hljs-title"># class TreeNode:</span> <span class="hljs-title"># def __init__(self, x):</span> <span class="hljs-title"># self.val = x</span> <span class="hljs-title"># self.left = None</span> <span class="hljs-title"># self.right = None</span> <span class="hljs-class"><span class="hljs-keyword">class</span> <span class="hljs-title">FindElements</span>:</span> <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">__init__</span><span class="hljs-params">(self, root: TreeNode)</span>:</span> <span class="hljs-title"># set 不能放在init外側。 因為測試用例之間不會銷毀FindElements的變量</span> self.seen = set() <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">recover</span><span class="hljs-params">(node)</span>:</span> <span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> node: <span class="hljs-keyword">return</span> node; <span class="hljs-keyword">if</span> node.left: node.left.val = <span class="hljs-params">2</span> * node.val + <span class="hljs-params">1</span> self.seen.add(node.left.val) <span class="hljs-keyword">if</span> node.right: node.right.val = <span class="hljs-params">2</span> * node.val + <span class="hljs-params">2</span> self.seen.add(node.right.val) recover(node.left) recover(node.right) <span class="hljs-keyword">return</span> node root.val = <span class="hljs-params">0</span> self.seen.add(<span class="hljs-params">0</span>) self.node = recover(root) <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">find</span><span class="hljs-params">(self, target: int)</span> -> bool:</span> <span class="hljs-keyword">return</span> target <span class="hljs-keyword">in</span> self.seen <span class="hljs-title"># Your FindElements object will be instantiated and called as such:</span> <span class="hljs-title"># obj = FindElements(root)</span> <span class="hljs-title"># param_1 = obj.find(target)</span> ``` ``` 這種解法可以 AC，但是在數據量非常大的時候，可能 MLE，我們繼續考慮優化。 ## 二進制法 ### 思路 這是一種非常巧妙的做法。 如果我們把樹中的數全部加 1 會怎么樣？ （圖參考 [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://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%EF%BC%89) 仔細觀察發現，每一行的左右子樹分別有不同的前綴：  Ok，那么算法就來了。為了便于理解，我們來舉個具體的例子，比如 target 是 9，我們首先將其加 1，二進制表示就是 1010。不考慮第一位，就是 010，我們只要： - 0 向左 ?? - 1 向右 ?? - - 0 向左 ?? 就可以找到 9 了。 > 0 表示向左 ， 1 表示向右 ### 代碼 ``` <pre class="calibre18">``` <span class="hljs-title"># Definition for a binary tree node.</span> <span class="hljs-title"># class TreeNode:</span> <span class="hljs-title"># def __init__(self, x):</span> <span class="hljs-title"># self.val = x</span> <span class="hljs-title"># self.left = None</span> <span class="hljs-title"># self.right = None</span> <span class="hljs-class"><span class="hljs-keyword">class</span> <span class="hljs-title">FindElements</span>:</span> node = <span class="hljs-keyword">None</span> <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">__init__</span><span class="hljs-params">(self, root: TreeNode)</span>:</span> <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">recover</span><span class="hljs-params">(node)</span>:</span> <span class="hljs-keyword">if</span> <span class="hljs-keyword">not</span> node: <span class="hljs-keyword">return</span> node; <span class="hljs-keyword">if</span> node.left: node.left.val = <span class="hljs-params">2</span> * node.val + <span class="hljs-params">1</span> <span class="hljs-keyword">if</span> node.right: node.right.val = <span class="hljs-params">2</span> * node.val + <span class="hljs-params">2</span> recover(node.left) recover(node.right) <span class="hljs-keyword">return</span> node root.val = <span class="hljs-params">0</span> self.node = recover(root) <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">find</span><span class="hljs-params">(self, target: int)</span> -> bool:</span> node = self.node <span class="hljs-keyword">for</span> bit <span class="hljs-keyword">in</span> bin(target+<span class="hljs-params">1</span>)[<span class="hljs-params">3</span>:]: node = node <span class="hljs-keyword">and</span> (node.left, node.right)[int(bit)] <span class="hljs-keyword">return</span> bool(node) <span class="hljs-title"># Your FindElements object will be instantiated and called as such:</span> <span class="hljs-title"># obj = FindElements(root)</span> <span class="hljs-title"># param_1 = obj.find(target)</span> ``` ``` ## 關鍵點解析 - 空間換時間 - 二進制思維 - 將 target + 1