bst.py

#!/usr/bin/env python


class TreeNode:

    def __init__(self):
        self.value = None
        self.left = None
        self.right = None
        self.size = 0

    def insert(self, value):
        if self.value == value:
            return

        self.size += 1
        if not self.value:
            self.value = value
            return

        if value < self.value:
            if not self.left:
                self.left = TreeNode()
            self.left.insert(value)
        else:
            if not self.right:
                self.right = TreeNode()
            self.right.insert(value)

    def traverse_in_order(self, visited):
        if self.left:
            self.left.traverse_in_order(visited)

        visited.append(self.value)

        if self.right:
            self.right.traverse_in_order(visited)

    def traverse_pre_order(self, visited):
        visited.append(self.value)

        if self.left:
            self.left.traverse_pre_order(visited)

        if self.right:
            self.right.traverse_pre_order(visited)

    def traverse_post_order(self, visited):
        if self.left:
            self.left.traverse_post_order(visited)

        if self.right:
            self.right.traverse_post_order(visited)

        visited.append(self.value)

    def traverse_level(self, visited):
        visited.append(self.value)
        queue = CustomQueue()
        queue.put(self)
        while not queue.empty():
            current_node = queue.push()
            if current_node.left:
                visited.append(current_node.left.value)
                queue.put(current_node.left)
            if current_node.right:
                visited.append(current_node.right.value)
                queue.put(current_node.right)

    def exists(self, value):
        if self.value == value:
            return True

        if value < self.value and self.left:
            return self.left.exists(value)
        elif self.right:
            return self.right.exists(value)

    def min(self):
        if self.left:
            return self.left.min()
        else:
            return self.value

    def max(self):
        if self.right:
            return self.right.max()
        else:
            return self.value

    def get_sum_by_level(self, level):
        return self._sum_by_level(self, 0, level)

    def _sum_by_level(self, node, depth, level):
        if not node or depth > level:
            return 0

        if depth == level:
            return node.value

        sum_ = 0
        if node.left:
            sum_ = self._sum_by_level(node.left, depth + 1, level)

        if node.right:
            sum_ += self._sum_by_level(node.right, depth + 1, level)

        return sum_


class BinarySearchTree:

    def __init__(self):
        self.root = None

    def insert(self, value):

        if not self.root:
            self.root = TreeNode()

        self.root.insert(value)

    def traverse_in_order(self):
        if not self.root:
            return []
        visited = []
        self.root.traverse_in_order(visited)
        return visited

    def traverse_pre_order(self):
        if not self.root:
            return []
        visited = []
        self.root.traverse_pre_order(visited)
        return visited

    def traverse_post_order(self):
        if not self.root:
            return []
        visited = []
        self.root.traverse_post_order(visited)
        return visited

    def traverse_level(self):
        if not self.root:
            return []
        visited = []
        self.root.traverse_level(visited)
        return visited

    def exists(self, value):
        if not self.root:
            return False
        return self.root.exists(value)

    def min(self):
        if self.root:
            return self.root.min()

    def max(self):
        if self.root:
            return self.root.max()

    def get_sum_by_level(self, level):
        if self.root:
            return self.root.get_sum_by_level(level)

    def size(self):
        return self.root.size


class CustomQueue:

    def __init__(self):
        self.queue = []

    def push(self):
        return self.queue.pop()

    def put(self, value):
        self.queue.insert(0, value)

    def empty(self):
        return len(self.queue) == 0

bst_test.py

#!/usr/bin/env python

import unittest

from bst import BinarySearchTree


class BinarySearchTreeTest(unittest.TestCase):

    def test_should_insert_element(self):
        tree = self.get_bst()
        actual = tree.size()
        expetected = 9
        self.assertEqual(expetected, actual)

    def test_should_traverse_in_order(self):
        tree = self.get_bst()
        expected = [15, 20, 22, 25, 26, 27, 29, 30, 32]
        actual = tree.traverse_in_order()
        self.assertEqual(expected, actual)

    def test_should_traverse_pre_order(self):
        tree = self.get_bst()
        expected = [25, 20, 15, 22, 27, 26, 30, 29, 32]
        actual = tree.traverse_pre_order()
        self.assertEqual(expected, actual)

    def test_should_traverse_post_order(self):
        tree = self.get_bst()
        expected = [15, 22, 20, 26, 29, 32, 30, 27, 25]
        actual = tree.traverse_post_order()
        self.assertEqual(expected, actual)

    def test_should_traverse_level(self):
        tree = self.get_bst()
        expected = [25, 20, 27, 15, 22, 26, 30, 29, 32]
        actual = tree.traverse_level()
        self.assertEqual(expected, actual)

    def test_should_check_if_exists(self):
        tree = self.get_bst()
        actual = tree.exists(15)
        self.assertTrue(actual)
        actual = tree.exists(150)
        self.assertFalse(actual)

    def test_should_find_min_value(self):
        tree = self.get_bst()
        expected = 15
        actual = tree.min()
        self.assertEqual(expected, actual)

    def test_should_find_max_value(self):
        tree = self.get_bst()
        expected = 32
        actual = tree.max()
        self.assertEqual(expected, actual)

    def test_should_get_sum_by_level(self):
        tree = self.get_bst()
        expected = 47
        actual = tree.get_sum_by_level(1)
        self.assertEqual(expected, actual)
        expected = 25
        actual = tree.get_sum_by_level(0)
        self.assertEqual(expected, actual)
        expected = 61
        actual = tree.get_sum_by_level(3)
        self.assertEqual(expected, actual)

    def get_bst(self):
        tree = BinarySearchTree()
        tree.insert(25)
        tree.insert(20)
        tree.insert(15)
        tree.insert(27)
        tree.insert(30)
        tree.insert(29)
        tree.insert(26)
        tree.insert(22)
        tree.insert(32)
        return tree


if __name__ == '__main__':
    unittest.main()