leetcode

Leetcode Solutions (0x6A73)
git clone https://git.pngu.org/leetcode
Log | Files | Refs | Submodules | README

contaminated_binary_tree.py (1985B)

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