about summary refs log tree commit diff
path: root/users/wpcarro/scratch/facebook/balanced-binary-tree.py
blob: afa9706f97ba9bd838efe8712c883399224cd00b (plain) (blame)
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
70
from collections import deque

class Node(object):
    # __init__ :: T(A)
    def __init__(self, value=None, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

    # insert_left :: T(A) -> A -> T(A)
    def insert_left(self, value):
        self.left = Node(value)
        return self.left

    # insert_right :: T(A) -> A -> T(A)
    def insert_right(self, value):
        self.right = Node(value)
        return self.right

    # is_superbalanced :: T(A) -> Bool
    def is_superbalanced(self):
        xs = deque()
        min_depth, max_depth = float('inf'), float('-inf')
        xs.append((self, 0))
        while xs:
            x, d = xs.popleft()
            # Only redefine the depths at leaf nodes
            if not x.left and not x.right:
                min_depth, max_depth = min(min_depth, d), max(max_depth, d)
            if x.left:
                xs.append((x.left, d + 1))
            if x.right:
                xs.append((x.right, d + 1))
        return max_depth - min_depth <= 1

    # __repr__ :: T(A) -> String
    def __repr__(self):
        result = ''
        xs = deque()
        xs.append((self, 0))
        while xs:
            node, indent = xs.popleft()
            result += '{i}{x}\n'.format(i=' ' * indent, x=node.value)
            if node.left:
                xs.append((node.left, indent + 2))
            if node.right:
                xs.append((node.right, indent + 2))
        return result

# from_array :: List(A) -> T(A)
def from_array(values):
    xs = deque()
    root = Node()
    xs.append(root)
    for value in values:
        node = xs.popleft()
        node.value = value
        node.left = Node()
        xs.append(node.left)
        node.right = Node()
        xs.append(node.right)
    return root

x = from_array([1, 1, 1, 1, 1, 1, 1])
print(x)
print(x.is_superbalanced())

x = Node(1, Node(2), Node(3))
print(x)
print(x.is_superbalanced())