Cracking_Coding_Interview/4.4.py at master · Abiaina/Cracking_Coding_Interview · 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
data = {
    "a": ["b","c"],
    "b": ["d"],
    "c": ["e", "f"],
    "d": [],
    "e": [],
    "f": []
}

data2 = {
    "a": ["b","c"],
    "b": [],
    "c": ["d"],
    "d": ["e"],
    "e": []
}

def check_bst_balance (data, root):
    if root == None or data[root][0] and data[root][1] == None:
        return True

    r_height = 0
    l_height = 0

    r_current = data[root][1]
    l_current = data[root][0]

    r_que = [r_current, None]
    l_que = [l_current, None]
    while len(r_que) > 0:
        r_current = r_que.pop(0)
        if r_current == None:
            r_height += 1
            if len(r_que) > 0:
                r_que.append(None)
            else:
                break
        else:
            for i in data[r_current]:
                r_que.append(i)

    while len(l_que) > 0:
        l_current = l_que.pop(0)
        if l_current == None:
            l_height += 1
            if len(l_que) > 0:
                l_que.append(None)
            else:
                break
        else:
            for i in data[l_current]:
                l_que.append(i)
    if (l_height - r_height) > 1 or (l_height - r_height) < -1:
        return False
    return True

if check_bst_balance(data, "a") and not check_bst_balance(data2, "a") :
    print("Passes!")