Problem of The Day: Find Elements in a Contaminated Binary Tree

Problem Statement

Intuition

When given a corrupted binary tree where every node’s value has been changed to -1, we need to recover it according to the given rules:

1. The root node’s value is set to 0.
2. For any node with value x, its left child gets the value 2 * x + 1, and its right child gets the value 2 * x + 2.

To efficiently check if a given value exists in the recovered tree, we can store the recovered values in a set.

Approach

1. Recovering the tree: We use a recursive function to traverse the given tree, assigning the correct values according to the rules and storing them in a set for quick lookup.
2. Finding an element: Since we stored all valid values in a set, checking for existence of a value is an O(1) operation.

Steps:

• Start at the root and assign it the value 0.
• Traverse the tree recursively:
  • Assign 2 * x + 1 to the left child.
  • Assign 2 * x + 2 to the right child.
• Store all valid values in a set.
• To find a target value, check if it exists in the set.

Complexity

• Time Complexity:

  • Recovering the tree takes \(O(n)\), where n is the number of nodes, as we visit each node once.
  • Searching for a target value takes \(O(1)\) using a set lookup.

• Space Complexity:

  • The set storing the elements takes \(O(n)\) space.
  • The recursion stack takes \(O(h)\) space in the worst case, where h is the height of the tree.
  • Overall, the worst-case space complexity is \(O(n)\).

Code

# 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.elements = set()
        self.root = self.recover(root, 0)

    def recover(self, root, x):
        new_node = TreeNode(x)
        self.elements.add(x)
        if root.left is not None:
            new_node.left = self.recover(root.left, 2 * x + 1)
        if root.right is not None:
            new_node.right = self.recover(root.right, 2 * x + 2)
        return new_node

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

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

Editorial

Approach 1: Tree Traversal (DFS)

class FindElements:
    def __init__(self, root: TreeNode):
        self.seen = set()
        self.dfs(root, 0)

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

    def dfs(self, current_node, current_value):
        if current_node is None:
            return
        # visit current node by adding its value to seen
        self.seen.add(current_value)
        self.dfs(current_node.left, current_value * 2 + 1)
        self.dfs(current_node.right, current_value * 2 + 2)

Approach 2: Tree Traversal (BFS)

class FindElements:

    def __init__(self, root: TreeNode):
        self.seen = set()
        self.bfs(root)

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

    def bfs(self, root: TreeNode) -> None:
        queue = [root]
        root.val = 0

        while queue:
            current_node = queue.pop(0)
            # visit current_node by adding its recovered value to the set
            self.seen.add(current_node.val)
            if current_node.left:
                current_node.left.val = current_node.val * 2 + 1
                queue.append(current_node.left)
            if current_node.right:
                current_node.right.val = current_node.val * 2 + 2
                queue.append(current_node.right)