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

impl Preorder for TreeLink {
    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 {
    root: TreeLink,
    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.