Skip to content

Commit 6937b4b

Browse files
authored
Update binary_search_tree.py
1 parent fe4aad0 commit 6937b4b

File tree

1 file changed

+63
-100
lines changed

1 file changed

+63
-100
lines changed

data_structures/binary_tree/binary_search_tree.py

+63-100
Original file line numberDiff line numberDiff line change
@@ -10,74 +10,65 @@
1010
/ \ /
1111
4 7 13
1212
13-
>>> t = BinarySearchTree()
14-
>>> t.insert(8, 3, 6, 1, 10, 14, 13, 4, 7)
15-
>>> print(" ".join(repr(i.value) for i in t.traversal_tree()))
16-
8 3 1 6 4 7 10 14 13
17-
18-
>>> tuple(i.value for i in t.traversal_tree(inorder))
13+
>>> tree = BinarySearchTree()
14+
>>> tree.insert(8, 3, 6, 1, 10, 14, 13, 4, 7)
15+
>>> tuple(node.value for node in tree.traversal_tree()) # inorder traversal (sorted)
1916
(1, 3, 4, 6, 7, 8, 10, 13, 14)
20-
>>> tuple(t)
17+
>>> tuple(node.value for node in tree.traversal_tree(postorder))
18+
(1, 4, 7, 6, 3, 13, 14, 10, 8)
19+
20+
>>> tuple(tree)
2121
(1, 3, 4, 6, 7, 8, 10, 13, 14)
22-
>>> t.find_kth_smallest(3, t.root)
23-
4
24-
>>> tuple(t)[3-1]
22+
>>> iter_t = iter(tree)
23+
>>> next(iter_t)
24+
1
25+
>>> next(iter_t)
26+
3
27+
>>> tuple(tree)[3-1] # 3rd smallest element in a zero-indexed tuple
2528
4
29+
>>> sum(tree)
30+
66
2631
27-
>>> print(" ".join(repr(i.value) for i in t.traversal_tree(postorder)))
28-
1 4 7 6 3 13 14 10 8
29-
>>> t.remove(20)
32+
>>> tuple(node.value for node in tree.traversal_tree(postorder))
33+
(1, 4, 7, 6, 3, 13, 14, 10, 8)
34+
>>> tree.remove(20)
3035
Traceback (most recent call last):
3136
...
3237
ValueError: Value 20 not found
33-
>>> BinarySearchTree().search(6)
34-
Traceback (most recent call last):
35-
...
36-
IndexError: Warning: Tree is empty! please use another.
3738
3839
Other example:
3940
40-
>>> testlist = (8, 3, 6, 1, 10, 14, 13, 4, 7)
41-
>>> t = BinarySearchTree()
42-
>>> for i in testlist:
43-
... t.insert(i)
41+
>>> values = (8, 3, 6, 1, 10, 14, 13, 4, 7)
42+
>>> tree = BinarySearchTree()
43+
>>> for value in values:
44+
... tree.insert(value)
4445
4546
Prints all the elements of the list in order traversal
46-
>>> print(t)
47+
>>> print(tree)
4748
{'8': ({'3': (1, {'6': (4, 7)})}, {'10': (None, {'14': (13, None)})})}
4849
4950
Test existence
50-
>>> t.search(6) is not None
51-
True
52-
>>> 6 in t
51+
>>> 6 in tree
5352
True
54-
>>> t.search(-1) is not None
55-
False
56-
>>> -1 in t
53+
>>> -1 in tree
5754
False
5855
59-
>>> t.search(6).is_right
56+
>>> tree.search(6).is_right
6057
True
61-
>>> t.search(1).is_right
58+
>>> tree.search(1).is_right
6259
False
6360
64-
>>> t.get_max().value
61+
>>> max(tree)
6562
14
66-
>>> max(t)
67-
14
68-
>>> t.get_min().value
69-
1
70-
>>> min(t)
63+
>>> min(tree)
7164
1
72-
>>> t.empty()
65+
>>> not tree
7366
False
74-
>>> not t
75-
False
76-
>>> for i in testlist:
77-
... t.remove(i)
78-
>>> t.empty()
79-
True
80-
>>> not t
67+
>>> for value in values:
68+
... tree.remove(value)
69+
>>> list(tree)
70+
[]
71+
>>> not tree
8172
True
8273
"""
8374
from __future__ import annotations
@@ -144,15 +135,12 @@ def __reassign_nodes(self, node: Node, new_children: Node | None) -> None:
144135
else:
145136
self.root = new_children
146137

147-
def empty(self) -> bool:
148-
return self.root is None
149-
150138
def __insert(self, value) -> None:
151139
"""
152140
Insert a new node in Binary Search Tree with value label
153141
"""
154142
new_node = Node(value) # create a new Node
155-
if self.empty(): # if Tree is empty
143+
if not self: # if Tree is empty
156144
self.root = new_node # set its root
157145
else: # Tree is not empty
158146
parent_node = self.root # from root
@@ -178,47 +166,32 @@ def insert(self, *values) -> None:
178166
self.__insert(value)
179167

180168
def search(self, value) -> Node | None:
181-
if self.empty():
169+
if not self:
182170
raise IndexError("Warning: Tree is empty! please use another.")
183-
else:
184-
node = self.root
185-
# use lazy evaluation here to avoid NoneType Attribute error
186-
while node is not None and node.value is not value:
187-
node = node.left if value < node.value else node.right
188-
return node
171+
node = self.root
172+
# use lazy evaluation here to avoid NoneType Attribute error
173+
while node and node.value is not value:
174+
node = node.left if value < node.value else node.right
175+
return node
189176

190177
def get_max(self, node: Node | None = None) -> Node | None:
191178
"""
192179
We go deep on the right branch
193180
"""
194181
if node is None:
195-
if self.root is None:
182+
if not self.root:
196183
return None
197184
node = self.root
198185

199-
if not self.empty():
186+
if self:
200187
while node.right is not None:
201188
node = node.right
202189
return node
203190

204-
def get_min(self, node: Node | None = None) -> Node | None:
205-
"""
206-
We go deep on the left branch
207-
"""
208-
if node is None:
209-
node = self.root
210-
if self.root is None:
211-
return None
212-
if not self.empty():
213-
node = self.root
214-
while node.left is not None:
215-
node = node.left
216-
return node
217-
218191
def remove(self, value: int) -> None:
219192
# Look for the node with that label
220193
node = self.search(value)
221-
if node is None:
194+
if not node:
222195
msg = f"Value {value} not found"
223196
raise ValueError(msg)
224197

@@ -229,29 +202,18 @@ def remove(self, value: int) -> None:
229202
elif node.right is None: # Has only left children
230203
self.__reassign_nodes(node, node.left)
231204
else:
232-
predecessor = self.get_max(
233-
node.left
234-
) # Gets the max value of the left branch
205+
# Gets the max value of the left branch
206+
predecessor = self.get_max(node.left)
235207
self.remove(predecessor.value) # type: ignore
236-
node.value = (
237-
predecessor.value # type: ignore
238-
) # Assigns the value to the node to delete and keep tree structure
208+
# Assigns the value to the node to delete and keep tree structure
209+
node.value = predecessor.value # type: ignore
239210

240-
def preorder_traverse(self, node: Node | None) -> Iterable:
241-
if node is not None:
211+
@classmethod
212+
def preorder_traverse(cls, node: Node | None) -> Iterable:
213+
if node:
242214
yield node # Preorder Traversal
243-
yield from self.preorder_traverse(node.left)
244-
yield from self.preorder_traverse(node.right)
245-
246-
def traversal_tree(self, traversal_function=None) -> Any:
247-
"""
248-
This function traversal the tree.
249-
You can pass a function to traversal the tree as needed by client code
250-
"""
251-
if traversal_function is None:
252-
return self.preorder_traverse(self.root)
253-
else:
254-
return traversal_function(self.root)
215+
yield from cls.preorder_traverse(node.left)
216+
yield from cls.preorder_traverse(node.right)
255217

256218
def inorder(self, arr: list, node: Node | None) -> None:
257219
"""Perform an inorder traversal and append values of the nodes to
@@ -261,29 +223,30 @@ def inorder(self, arr: list, node: Node | None) -> None:
261223
arr.append(node.value)
262224
self.inorder(arr, node.right)
263225

264-
def find_kth_smallest(self, k: int, node: Node) -> int:
265-
"""Return the kth smallest element in a binary search tree"""
266-
arr: list[int] = []
267-
self.inorder(arr, node) # append all values to list using inorder traversal
268-
return arr[k - 1]
226+
def traversal_tree(self, traversal_function=None) -> Any:
227+
"""
228+
This function traversal the tree.
229+
You can pass a function to traversal the tree as needed by client code
230+
"""
231+
return (traversal_function or inorder)(self.root)
269232

270233

271234
def inorder(curr_node: Node | None) -> list[Node]:
272235
"""
273236
inorder (left, self, right)
274237
"""
275238
node_list = []
276-
if curr_node is not None:
239+
if curr_node:
277240
node_list = inorder(curr_node.left) + [curr_node] + inorder(curr_node.right)
278241
return node_list
279242

280243

281244
def postorder(curr_node: Node | None) -> list[Node]:
282245
"""
283-
postOrder (left, right, self)
246+
postorder (left, right, self)
284247
"""
285248
node_list = []
286-
if curr_node is not None:
249+
if curr_node:
287250
node_list = postorder(curr_node.left) + postorder(curr_node.right) + [curr_node]
288251
return node_list
289252

0 commit comments

Comments
 (0)