32 lines
838 B
Python
32 lines
838 B
Python
# https://leetcode.com/problems/find-elements-in-a-contaminated-binary-tree
|
|
|
|
from typing import Optional, Set
|
|
|
|
class TreeNode:
|
|
def __init__(self, val=0, left=None, right=None):
|
|
self.val = val
|
|
self.left = left
|
|
self.right = right
|
|
|
|
class FindElements:
|
|
r: Optional[TreeNode]
|
|
s: Set[int]
|
|
|
|
def __init__(self, root: Optional[TreeNode]):
|
|
self.r = root
|
|
self.s = set()
|
|
if self.r is not None:
|
|
self.s.add(0)
|
|
self.recurse_init(self.r, 0)
|
|
|
|
def recurse_init(self, curr: TreeNode, curr_val: int):
|
|
if curr.left is not None:
|
|
self.s.add(curr_val * 2 + 1)
|
|
self.recurse_init(curr.left, curr_val * 2 + 1)
|
|
if curr.right is not None:
|
|
self.s.add(curr_val * 2 + 2)
|
|
self.recurse_init(curr.right, curr_val * 2 + 2)
|
|
|
|
def find(self, target: int) -> bool:
|
|
return target in self.s
|
|
|