diff --git a/01_binary_search_trees/bst.py b/01_binary_search_trees/bst.py index ed442dc..587aa47 100644 --- a/01_binary_search_trees/bst.py +++ b/01_binary_search_trees/bst.py @@ -130,6 +130,99 @@ class BinarySearchTree(): return (left and right) + +class BinarySearchIterator(): + """ + Transforms a binary search tree into an iterator. + Will returns the values of the tree in ascending order. + """ + + def __init__(self, tree): + self.tree = tree + self.node = tree.root_node + self.max_ite = 1000000000 # max number of nodes + self.ite = 0 + self.pos = "left" + + def __iter__(self): + return self + + def go_left(self): + while(self.node.has_left_child()): + self.node = self.node.left_child + + ret_val = self.node.value + + self.pos = "top" # next state is top + + return ret_val + + def go_top(self): + self.node = self.node.parent + ret_val = self.node.value + + self.pos = "right" + + return ret_val + + def go_right(self): + if (not self.node.has_right_child()): + raise StopIteration # end condition. We reached the end of tree + else: + self.node = self.node.right_child + if self.node.has_left_child(): + self.pos = "left" + else: + self.pos = "right" + + ret_val = self.node.value + return ret_val # we dont change node here + + def next(self): + if self.ite > self.max_ite: + print "max limit reached!" + raise StopIteration # max number of elements + + print self.pos + if self.pos == "left": + return self.go_left() + elif self.pos == "top": + return self.go_top() + elif self.pos == "right": + try: + return self.go_right() + except StopIteration: + raise StopIteration + else: + print "problem!" + raise StopIteration # problem somewhere + + self.ite += 1 + + # max_iter = 1000000000 # max number of nodes + # ite = 0 + # ret = "" + # node = self.root_node + + # top = False # dictates when to climb one level + # while(ite < max_iter): + # while (node.has_left_child()): + # node = node.left_child + # top = True + # ret += str(node.value) + " " + # if top: # we need to climb one level + # node = node.parent + # top = False + # ret += str(node.value) + " " + # if (not node.has_right_child()): + # return ret.strip() # removes superflous spaces + # else: + # node = node.right_child + # ite += 1 + # return None # problem + + + class BinarySearchNode(): """ Defines any node of the Binary Search Tree. @@ -162,4 +255,4 @@ class BinarySearchNode(): if self.has_right_child(): ret += " " + str(self.right_child.value) - return ret \ No newline at end of file + return ret diff --git a/01_binary_search_trees/bst_test.py b/01_binary_search_trees/bst_test.py index 101a8a1..8813913 100644 --- a/01_binary_search_trees/bst_test.py +++ b/01_binary_search_trees/bst_test.py @@ -5,6 +5,7 @@ Unit tests for the binary search tree data structure """ from bst import BinarySearchTree +from bst import BinarySearchIterator import unittest @@ -234,5 +235,35 @@ class test_binary_search_tree(unittest.TestCase): self.assertEqual(bst.__str__(), "4 6 10 16") + def test_iterator(self): + + bst = BinarySearchTree() + + self.assertEqual(bst.min(), None) + + node_val1 = 6 + bst.add(node_val1) + node_val2 = 4 + bst.add(node_val2) + node_val3 = 10 + bst.add(node_val3) + node_val4 = 16 + bst.add(node_val4) + + # 6 + # / \ + # 4 10 + # \ + # 16 + + ite = BinarySearchIterator(bst) + self.assertEquals(ite.tree, bst) + print ite.next() + print ite.next() + print ite.next() + print ite.next() + + + if __name__ == '__main__': unittest.main()