1261. Find Elements in a Contaminated Binary Tree

Given a binary tree with the following rules:

1. root.val == 0
2. If treeNode.val == x and treeNode.left != null, then treeNode.left.val == 2 * x + 1
3. If treeNode.val == x and treeNode.right != null, then treeNode.right.val == 2 * x + 2

Now the binary tree is contaminated, which means all treeNode.val have been changed to -1.

You need to first recover the binary tree and then implement the FindElements class:

• FindElements(TreeNode* root) Initializes the object with a contamined binary tree, you need to recover it first.
• bool find(int target) Return if the target value exists in the recovered binary tree.

Example 1:

Input
["FindElements","find","find"]
[[[-1,null,-1]],[1],[2]]
Output
[null,false,true]
Explanation
FindElements findElements = new FindElements([-1,null,-1]);
findElements.find(1); // return False
findElements.find(2); // return True

Example 2:

Input
["FindElements","find","find","find"]
[[[-1,-1,-1,-1,-1]],[1],[3],[5]]
Output
[null,true,true,false]
Explanation
FindElements findElements = new FindElements([-1,-1,-1,-1,-1]);
findElements.find(1); // return True
findElements.find(3); // return True
findElements.find(5); // return False

Example 3:

Input
["FindElements","find","find","find","find"]
[[[-1,null,-1,-1,null,-1]],[2],[3],[4],[5]]
Output
[null,true,false,false,true]
Explanation
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

Constraints:

• TreeNode.val == -1
• The height of the binary tree is less than or equal to 20
• The total number of nodes is between [1, 10^4]
• Total calls of find() is between [1, 10^4]
• 0 <= target <= 10^6

Rust Solution

use rustgym_util::*;
use std::collections::HashSet;

trait Preorder {
fn recover(&mut self, x: i32, hs: &mut HashSet<i32>);
}

fn recover(&mut self, x: i32, hs: &mut HashSet<i32>) {
if let Some(node) = self {
hs.insert(x);
node.borrow_mut().val = x;
node.borrow_mut().left.recover(2 * x + 1, hs);
node.borrow_mut().right.recover(2 * x + 2, hs);
}
}
}

struct FindElements {
hs: HashSet<i32>,
}

impl FindElements {
fn new(mut root: TreeLink) -> Self {
let mut hs = HashSet::new();
root.recover(0, &mut hs);
FindElements { root, hs }
}

fn find(&self, target: i32) -> bool {
self.hs.contains(&target)
}
}

#[test]
fn test() {
let root = tree!(-1, None, tree!(-1));
let fe = FindElements::new(root);
let target = 1;
let res = false;
assert_eq!(fe.find(target), res);
let target = 2;
let res = true;
assert_eq!(fe.find(target), res);
let root = tree!(-1, tree!(-1, tree!(-1), tree!(-1)), tree!(-1));
let fe = FindElements::new(root);
let target = 1;
let res = true;
assert_eq!(fe.find(target), res);
let target = 3;
let res = true;
assert_eq!(fe.find(target), res);
let target = 5;
let res = false;
assert_eq!(fe.find(target), res);
let root = tree!(-1, None, tree!(-1, tree!(-1, tree!(-1), None), None));
let fe = FindElements::new(root);
let target = 2;
let res = true;
assert_eq!(fe.find(target), res);
let target = 3;
let res = false;
assert_eq!(fe.find(target), res);
let target = 4;
let res = false;
assert_eq!(fe.find(target), res);
let target = 5;
let res = true;
assert_eq!(fe.find(target), res);
}

Having problems with this solution? Click here to submit an issue on github.