Algorithms/trees.py at master · qedpi/Algorithms · GitHub

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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
from collections import defaultdict

# Binary-Tree Node
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None


x = TreeNode(3); x.left = TreeNode(1); x.left.left = TreeNode(0); x.right = TreeNode(4)


def traverse_in_order(root):
    if root.left:
        yield from traverse_in_order(root.left)
    yield root.val
    if root.right:
        yield from traverse_in_order(root.right)
#print(*traverse_in_order(x))


def traverse_pre_order(root):
    yield root.val
    if root.left:
        yield from traverse_pre_order(root.left)
    if root.right:
        yield from traverse_pre_order(root.right)
#print(*traverse_pre_order(x))


def traverse_post_order(root):
    if root.left:
        yield from traverse_post_order(root.left)
    if root.right:
        yield from traverse_post_order(root.right)
    yield root.val
#print(*traverse_post_order(x))


def traverse_in_order_it(root):
    q = [root]
    while q:
        node = q.pop()
        if node:
            if type(node) is tuple:
                yield node[0]
            else:
                q.extend([node.right, (node.val,), node.left])
#print(*traverse_in_order_it(x))


def traverse_pre_order_it(root):
    q = [root]
    while q:
        node = q.pop()
        if node:
            yield node.val
            q.extend([node.right, node.left])
#print(*traverse_pre_order_it(x))


def traverse_post_order_it(root):
    q = [root]
    while q:
        node = q.pop()
        if node:
            if type(node) is tuple:
                yield node[0]
            else:
                q.extend([(node.val,), node.right, node.left])
#print(*traverse_post_order_it(x))


def leaves_values(root):
    if not root:
        return
    else:
        if not (root.left or root.right):
            yield root.val
        else:
            yield from leaves_values(root.left)
            yield from leaves_values(root.right)
#print(*leaves_values(x))


def tree_height(root):
    if not root:
        return 0
    else:
        return 1 + max(map(tree_height, (root.left, root.right)))
#print(tree_height(x))


def tree_layers_rel(root):
    layers = defaultdict(list)

    def traverse(root):
        if not root:
            return 0
        else:
            height = 1 + max(map(traverse, (root.left, root.right)))
            layers[height].append(root.val)
            return height

    traverse(root)
    return [(h + 1, layers[h + 1]) for h in range(len(layers))]
#print(*tree_layers_rel(x))


def tree_layers(root):
    layers = defaultdict(list)

    def traverse(root, h=0):
        if root:
            layers[h].append(root.val)
            [traverse(r, h + 1) for r in (root.left, root.right)]
        else:
            return

    traverse(root)

    return [(h + 1, layers[h]) for h in range(len(layers))]
#print(*tree_layers(x))


# aka invert
def reflect_tree(root):
    if root:
        root.left, root.right = map(reflect_tree, (root.right, root.left))
    return root
#print(*traverse_in_order(reflect_tree(x)))


def tree_equals(ta, tb):
    if (ta and not tb) or (tb and not ta):  # only one of them has children
        return False
    else:
        return not (ta or tb) or\
               (ta.val == tb.val and tree_equals(ta.left, tb.left) and tree_equals(ta.right, tb.right))
#print(tree_equals(x, x))


def _bst_balanced(root):
    if not root:
        return 1
    else:
        left, right = map(_bst_balanced, (root.left, root.right))
        return max(left, right) + 1 if all((left, right)) and abs(left - right) <= 1 else 0


def bst_balanced(root):
    return bool(_bst_balanced(root))

#x.left.left.left = TreeNode(5)
#print(bst_balanced(x))