Code4Fun/tree-orders.py at master · sban2009/Code4Fun · 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
"""
Given a tree, print its
1. maximum width (i.e. maximum nodes at any particular level)
2. level order traversal (BFS or level by level traversal)
3. spiral/zigzag traversal (similar to level order)
"""
from collections import deque


class node:
    def __init__(self, data) -> None:
        self.data = data
        self.left = None
        self.right = None


def max_width(root):
    # returns maximum width (max. #nodes at a level) in BT (5,15,7=> 3 nodes)
    if not root:
        return 0
    q = deque()
    q.append(root)
    maxsize = 0
    while q:
        count = len(q)
        maxsize = max(count, maxsize)
        while count:
            x = q.popleft()
            # print(x.data, q)
            count -= 1
            if x.left is not None:
                q.append(x.left)
            if x.right is not None:
                q.append(x.right)
    return maxsize


def level_order(root):
    # returns list of list containing nodes of BT, level-wise (3,9,20,5,15,7)
    if not root:
        return 0
    q = deque()
    q.append(root)
    result = []
    while q:
        count = len(q)
        result.append([i.data for i in q])
        while count:
            x = q.popleft()
            # print(x.data, q)
            count -= 1
            if x.left is not None:
                q.append(x.left)
            if x.right is not None:
                q.append(x.right)
    return result


def spiral(root):
    # returns zigzag/spiral order of nodes (3,20,9,5,15,7)
    if not root:
        return 0
    q = deque()
    q.append(root)
    level = -1  # *changed code from level_order
    result = []
    while q:
        count = len(q)
        level += 1  # *changed code from level_order
        if level % 2 == 0:  # *changed code from level_order
            result.append([i.data for i in q])
        else:
            result.append([i.data for i in reversed(q)])
        while count:
            x = q.popleft()
            # print(x.data, q)
            count -= 1
            if x.left is not None:
                q.append(x.left)
            if x.right is not None:
                q.append(x.right)
    return result


def vertical_order(root, hd, m):
    if not root:
        return
    try:
        m[hd].append(root.data)
    except:
        m[hd] = [root.data]
    vertical_order(root.left, hd - 1, m)
    vertical_order(root.right, hd + 1, m)


def vertical_order_print(root):
    hm = dict()
    vertical_order(root, 0, hm)
    for hd in sorted(hm):
        print(hm[hd], end=" ")


def main():
    """
    binary tree:
          3
       /     \
      9       20
     /      /   \
    5     15     7
    """
    root = node(3)
    root.left = node(9)
    root.right = node(20)
    root.left.left = node(5)
    root.right.left = node(15)
    root.right.right = node(7)

    print(max_width(root))
    print(level_order(root))
    print(spiral(root))
    vertical_order_print(root)


if __name__ == "__main__":
    main()