blob: 03de0350d898ee14c72c5a088b4fb6ad4dcfa150 (
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
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
|
import unittest
from collections import deque
# is_balanced :: Node(a) -> Bool
def is_balanced(node):
q = deque()
q.append((0, node))
mn, mx = None, None
while q:
depth, node = q.popleft()
# Current node is a leaf node
if not node.left and not node.right:
mx = depth if mx is None else max(mx, depth)
mn = depth if mn is None else min(mn, depth)
if mx - mn > 1:
return False
if node.left:
q.append((depth + 1, node.left))
if node.right:
q.append((depth + 1, node.right))
return mx - mn <= 1
# Tests
class Test(unittest.TestCase):
class BinaryTreeNode(object):
def __init__(self, value):
self.value = value
self.left = None
self.right = None
def insert_left(self, value):
self.left = Test.BinaryTreeNode(value)
return self.left
def insert_right(self, value):
self.right = Test.BinaryTreeNode(value)
return self.right
def test_full_tree(self):
tree = Test.BinaryTreeNode(5)
left = tree.insert_left(8)
right = tree.insert_right(6)
left.insert_left(1)
left.insert_right(2)
right.insert_left(3)
right.insert_right(4)
result = is_balanced(tree)
self.assertTrue(result)
def test_both_leaves_at_the_same_depth(self):
tree = Test.BinaryTreeNode(3)
left = tree.insert_left(4)
right = tree.insert_right(2)
left.insert_left(1)
right.insert_right(9)
result = is_balanced(tree)
self.assertTrue(result)
def test_leaf_heights_differ_by_one(self):
tree = Test.BinaryTreeNode(6)
left = tree.insert_left(1)
right = tree.insert_right(0)
right.insert_right(7)
result = is_balanced(tree)
self.assertTrue(result)
def test_leaf_heights_differ_by_two(self):
tree = Test.BinaryTreeNode(6)
left = tree.insert_left(1)
right = tree.insert_right(0)
right_right = right.insert_right(7)
right_right.insert_right(8)
result = is_balanced(tree)
self.assertFalse(result)
def test_three_leaves_total(self):
tree = Test.BinaryTreeNode(1)
left = tree.insert_left(5)
right = tree.insert_right(9)
right.insert_left(8)
right.insert_right(5)
result = is_balanced(tree)
self.assertTrue(result)
def test_both_subtrees_superbalanced(self):
tree = Test.BinaryTreeNode(1)
left = tree.insert_left(5)
right = tree.insert_right(9)
right_left = right.insert_left(8)
right.insert_right(5)
right_left.insert_left(7)
result = is_balanced(tree)
self.assertFalse(result)
def test_both_subtrees_superbalanced_two(self):
tree = Test.BinaryTreeNode(1)
left = tree.insert_left(2)
right = tree.insert_right(4)
left.insert_left(3)
left_right = left.insert_right(7)
left_right.insert_right(8)
right_right = right.insert_right(5)
right_right_right = right_right.insert_right(6)
right_right_right.insert_right(9)
result = is_balanced(tree)
self.assertFalse(result)
def test_only_one_node(self):
tree = Test.BinaryTreeNode(1)
result = is_balanced(tree)
self.assertTrue(result)
def test_linked_list_tree(self):
tree = Test.BinaryTreeNode(1)
right = tree.insert_right(2)
right_right = right.insert_right(3)
right_right.insert_right(4)
result = is_balanced(tree)
self.assertTrue(result)
unittest.main(verbosity=2)
|
