blob: 25cc339de888b2b614d760ff85cb59c129e8007b (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
|
# 1261. Find Elements in a Contaminated Binary Tree
"""
given a binary tree with the following rules: 'root.val == 0', for any
'TreeNode': if 'TreeNode.val' has a value 'x' and 'TreeNode.left != null'
then 'TreeNode.left.val == 2 x + 1', and if 'TreeNode.val' has a value '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. implemented the 'FindElements' class: 'findElementsCreate':
initialises the object with a contaminated tree and recovers it. 'bool
find(int target)' returns 'true' if the 'target' value exists in the
recovered binary tree.
"""
# Definition for a binary tree node.
class TreeNode(object):
def __init__(self, val=0, left=None, right=None):
self.val = val
self.left = left
self.right = right
class FindElements(object):
def __init__(self, root):
"""
:type root: Optional[TreeNode]
"""
def recover(root, val):
return (
root
and {val}
| recover(root.left, 2 * val + 1)
| recover(root.right, 2 * val + 2)
or set()
)
self.find = recover(root, 0).__contains__
if __name__ == "__main__":
r1 = TreeNode(-1)
r1.left = None
r1.right = TreeNode(-1)
r2 = TreeNode(-1)
r2.left = TreeNode(-1)
r2.right = TreeNode(-1)
r2.left.left = TreeNode(-1)
r2.left.right = TreeNode(-1)
r3 = TreeNode(-1)
r3.left = None
r3.right = TreeNode(-1)
r3.right.left = TreeNode(-1)
r3.right.right = None
r3.right.left.left = TreeNode(-1)
r3.right.left.right = None
o1 = FindElements(r1)
o2 = FindElements(r2)
o3 = FindElements(r3)
print(o1.find(1))
print(o1.find(2))
print(o2.find(1))
print(o2.find(3))
print(o2.find(5))
print(o3.find(2))
print(o3.find(3))
print(o3.find(4))
print(o3.find(5))
|