LeetCode 1261 - Find Elements in a Contaminated Binary Tree

Difficulty: 🟡 Medium
Topics: Hash Table, Tree, Depth-First Search, Breadth-First Search, Design, Binary Tree

Solution

Problem Understanding

The problem gives us a binary tree where every node value has been contaminated and replaced with -1. However, we know the exact rule that originally generated the values in the tree.

The original tree followed these rules:

The root node always had value 0
If a node had value x:
Its left child had value 2 * x + 1
Its right child had value 2 * x + 2

We are asked to design a class called FindElements that restores the tree and efficiently answers queries about whether a target value exists in the recovered tree.

The constructor receives the contaminated root node. During initialization, we must recover the original values. Then, the find(target) method should return True if the recovered tree contains target, otherwise False.

The constraints are important because they guide the design:

The tree can contain up to 10^4 nodes
There can also be up to 10^4 calls to find()
Target values can be as large as 10^6

A naive search through the tree for every query would be too expensive in the worst case. Since many find() calls may occur, preprocessing the tree once during construction is the key optimization.

Several edge cases are important:

A tree may contain only the root node
The tree may be heavily skewed to one side
Some nodes may be missing, meaning not every mathematically possible value actually exists
Queries may ask for values far larger than any value in the tree
The height can reach 20, which means recursive DFS is still safe in Python and Go

Approaches

Brute Force Approach

A straightforward solution is to recover the tree during construction and then perform a tree traversal every time find(target) is called.

The constructor would restore all node values using DFS or BFS. Then, each find() operation would traverse the tree searching for the target value.

This approach is correct because the recovery rules uniquely determine every node value, and a traversal eventually checks every node.

However, this becomes inefficient when many queries are made. If the tree contains n nodes and we perform q queries, the total complexity becomes O(n * q) in the worst case. With both values reaching 10^4, this could become too slow.

Optimal Approach

The key observation is that after recovering the tree, we repeatedly need fast existence checks.

A hash set is ideal for this scenario.

While recovering the tree, we store every recovered value inside a set. Then, each find(target) operation becomes a constant time hash lookup.

The recovery traversal itself only needs to visit each node once, so preprocessing costs O(n). After that, every query is O(1) on average.

This transforms the repeated expensive searches into fast membership checks.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n) per find	O(h) recursion stack	Traverses tree for every query
Optimal	O(n) preprocessing + O(1) per find	O(n)	Uses hash set for fast lookup

Algorithm Walkthrough

Initialize an empty hash set to store all recovered node values.
Set the root value to 0 because the problem guarantees the original root value was always 0.
Start a DFS traversal from the root.
For every visited node:

Insert its recovered value into the hash set
If the node has a left child:
Recover the left child value using 2 * current + 1
Continue DFS on the left child
If the node has a right child:
Recover the right child value using 2 * current + 2
Continue DFS on the right child

After preprocessing completes, the set contains every valid value in the tree.
For each find(target) call:

Return whether target exists in the hash set

The hash set is chosen because it provides average O(1) membership testing, which is ideal for repeated queries.

Why it works

The recovery rules uniquely determine every node value from its parent value. Since the root is known to be 0, recursively applying the left and right formulas reconstructs the exact original tree values.

The DFS visits every existing node exactly once, so every valid recovered value is inserted into the set. Therefore, the set precisely represents the contents of the recovered tree. Any target exists in the tree if and only if it exists in the set.

Python Solution

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right

from typing import Optional

class FindElements:

    def __init__(self, root: Optional[TreeNode]):
        self.values = set()
        self.root = root

        if self.root:
            self._recover(self.root, 0)

    def _recover(self, node: Optional[TreeNode], value: int) -> None:
        if not node:
            return

        node.val = value
        self.values.add(value)

        if node.left:
            self._recover(node.left, 2 * value + 1)

        if node.right:
            self._recover(node.right, 2 * value + 2)

    def find(self, target: int) -> bool:
        return target in self.values

# Your FindElements object will be instantiated and called as such:
# obj = FindElements(root)
# param_1 = obj.find(target)

The implementation begins by creating a hash set called values. This set stores every recovered node value for fast lookup.

The constructor saves the root reference and starts recovery from value 0. The helper method _recover() performs DFS traversal while assigning correct values.

Inside _recover(), the current node value is restored and inserted into the set. Then, the algorithm recursively computes child values using the problem's formulas.

The find() method becomes extremely simple because all preprocessing has already been completed. It only checks whether the target exists in the set.

Go Solution

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */

type FindElements struct {
    values map[int]bool
}

func Constructor(root *TreeNode) FindElements {
    obj := FindElements{
        values: make(map[int]bool),
    }

    if root != nil {
        obj.recover(root, 0)
    }

    return obj
}

func (this *FindElements) recover(node *TreeNode, value int) {
    if node == nil {
        return
    }

    node.Val = value
    this.values[value] = true

    if node.Left != nil {
        this.recover(node.Left, 2*value+1)
    }

    if node.Right != nil {
        this.recover(node.Right, 2*value+2)
    }
}

func (this *FindElements) Find(target int) bool {
    return this.values[target]
}

/**
 * Your FindElements object will be instantiated and called as such:
 * obj := Constructor(root);
 * param_1 := obj.Find(target);
 */

The Go implementation closely mirrors the Python version.

Instead of Python's set, Go uses map[int]bool to simulate a hash set. Presence in the map indicates that the value exists in the recovered tree.

Go requires explicit nil checks before recursion. The recursive helper method is attached to the FindElements struct.

Integer overflow is not an issue because the constraints keep values safely within Go's standard integer range.

Worked Examples

Example 1

Input:

[-1, null, -1]

Tree structure:

    -1
      \
      -1

Recovery process:

Step	Current Node	Assigned Value	Set Contents
1	Root	0	{0}
2	Right child	2	{0, 2}

Now evaluate queries:

Query	Exists?	Result
1	No	False
2	Yes	True

Example 2

Input:

[-1,-1,-1,-1,-1]

Tree structure:

Recovery traversal:

Step	Node Value	Recovered Value	Set Contents
1	Root	0	{0}
2	Left child	1	{0,1}
3	Left-left	3	{0,1,3}
4	Left-right	4	{0,1,3,4}
5	Right child	2	{0,1,2,3,4}

Queries:

Query	Exists?
1	True
3	True
5	False

Example 3

Input:

[-1,null,-1,-1,null,-1]

Tree structure:

Recovery:

Step	Node	Value
1	Root	0
2	Right child	2
3	Left child of 2	5
4	Left child of 5	11

Recovered values:

{0, 2, 5, 11}

Queries:

Query	Result
2	True
3	False
4	False
5	True

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n) preprocessing, O(1) per find	DFS visits each node once, set lookup is constant time
Space	O(n)	Hash set stores every recovered value

The constructor performs one DFS traversal across all nodes, so recovery costs O(n) time. Every node value is inserted once into the hash set.

Each find() call uses hash table membership testing, which is O(1) on average.

The space complexity is O(n) because every recovered value is stored inside the hash set. The recursion stack additionally uses O(h) space, where h is the tree height.

Test Cases

# Definition for testing
class TreeNode:
    def __init__(self, val=-1, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

# Example 1
root1 = TreeNode(-1, None, TreeNode(-1))
obj1 = FindElements(root1)

assert obj1.find(1) == False  # Missing left child value
assert obj1.find(2) == True   # Existing right child value

# Example 2
root2 = TreeNode(
    -1,
    TreeNode(-1, TreeNode(-1), TreeNode(-1)),
    TreeNode(-1)
)

obj2 = FindElements(root2)

assert obj2.find(1) == True   # Left child exists
assert obj2.find(3) == True   # Left-left child exists
assert obj2.find(4) == True   # Left-right child exists
assert obj2.find(5) == False  # Missing node

# Example 3
root3 = TreeNode(
    -1,
    None,
    TreeNode(
        -1,
        TreeNode(-1, TreeNode(-1)),
        None
    )
)

obj3 = FindElements(root3)

assert obj3.find(2) == True    # Right child exists
assert obj3.find(5) == True    # Recovered descendant
assert obj3.find(11) == True   # Deep descendant
assert obj3.find(3) == False   # Missing node

# Single node tree
root4 = TreeNode(-1)
obj4 = FindElements(root4)

assert obj4.find(0) == True   # Root always becomes 0
assert obj4.find(1) == False  # No children

# Left skewed tree
root5 = TreeNode(
    -1,
    TreeNode(
        -1,
        TreeNode(-1)
    )
)

obj5 = FindElements(root5)

assert obj5.find(1) == True   # Left child
assert obj5.find(3) == True   # Left-left child
assert obj5.find(2) == False  # Right child missing

# Large missing target
assert obj5.find(1000000) == False  # Very large absent value

Test	Why
Example 1	Verifies recovery of right child only
Example 2	Tests balanced tree recovery
Example 3	Tests sparse tree handling
Single node tree	Validates minimum input size
Left skewed tree	Ensures recursion handles uneven trees
Large missing target	Confirms absent values return false

Edge Cases

Tree Contains Only the Root

The smallest possible tree contains exactly one node. Since the root must recover to value 0, the only valid target is 0.

A buggy implementation might forget to initialize the root correctly or mishandle empty children. This solution explicitly starts recovery from value 0, so the root is always recovered correctly.

Sparse Trees With Missing Children

The tree may contain gaps where one child exists and the other does not. This is important because mathematically possible values are not guaranteed to exist.

For example, a node may have a right child without a left child. The implementation only recovers existing nodes by checking if node.left and if node.right, so nonexistent nodes are never incorrectly inserted into the set.

Very Large Target Queries

Queries may ask for values much larger than anything present in the tree.

A naive traversal based solution would still scan the entire tree for every such query. The hash set approach avoids this inefficiency because membership testing immediately returns False when the value is absent.