binary search tree


See On Github

Data

Contributor

Generic placeholder thumbnail

by dalleng

in python

Tags

search, tree

Source Code


class TreeNode:
    def __init__(self, value):
        self.value = value
        self.right = None
        self.left = None


class BinarySearchTree:
    def __init__(self):
        self.root = None

    def insert(self, value):
        if self.root is None:
            self.root = TreeNode(value)
        else:
            current = self.root
            while True:
                if value <= current.value:
                    if current.left:
                        current = current.left
                    else:
                        current.left = TreeNode(value)
                        break
                else:
                    if current.right:
                        current = current.right
                    else:
                        current.right = TreeNode(value)
                        break

    @staticmethod
    def find_min(node):
        if not node:
            return None
        else:
            current = node
            parent = None
            while True:
                if current.left:
                    parent = current
                    current = current.left
                else:
                    break
            return current, parent

    def delete_node(self, node, parent):
        if node.right and node.left:
            min_node, min_parent = self.find_min(node.right)
            node.value = min_node.value
            if min_parent is None:
                min_parent = node
            self.delete_node(min_node, min_parent)
        elif node.left:
            if parent.right == node:
                parent.right = node.left
            else:
                parent.left = node.left
        elif node.right:
            if parent.right == node:
                parent.right = node.right
            else:
                parent.left = node.right
        else:
            if parent.right == node:
                parent.right = None
            else:
                parent.left = None

    def delete(self, value):
        if not self.root:
            return None
        else:
            current = self.root
            parent = None
            while True:
                if value == current.value:
                    break
                elif value <= current.value:
                    if current.left:
                        parent = current
                        current = current.left
                    else:
                        return
                else:
                    if current.right:
                        parent = current
                        current = current.right
                    else:
                        return

            self.delete_node(current, parent)

    def search(self, value):
        if not self.root:
            return None
        else:
            current = self.root
            while True:
                if value == current.value:
                    return current
                elif value <= current.value:
                    if current.left:
                        current = current.left
                    else:
                        return None
                else:
                    if current.right:
                        current = current.right
                    else:
                        return None
import unittest
from bst import BinarySearchTree

class BinarySearchTreeTest(unittest.TestCase):
    def test_insertion(self):
        tree = BinarySearchTree()
        tree.insert(8)
        tree.insert(3)
        tree.insert(10)
        tree.insert(1)
        tree.insert(6)
        tree.insert(4)
        tree.insert(7)
        tree.insert(14)
        tree.insert(13)
        self.assertEquals(tree.root.value, 8)
        self.assertEquals(tree.root.left.value, 3)
        self.assertEquals(tree.root.left.left.value, 1)
        self.assertEquals(tree.root.left.right.value, 6)
        self.assertEquals(tree.root.left.right.left.value, 4)
        self.assertEquals(tree.root.left.right.right.value, 7)
        self.assertEquals(tree.root.right.value, 10)
        self.assertEquals(tree.root.right.right.value, 14)
        self.assertEquals(tree.root.right.right.left.value, 13)

    def test_search(self):
        # Emptry Tree
        self.assertEquals(BinarySearchTree().search(8), None)

        tree = BinarySearchTree()
        tree.insert(8)
        tree.insert(3)
        tree.insert(10)
        tree.insert(1)
        tree.insert(6)
        tree.insert(4)
        tree.insert(7)
        tree.insert(14)
        tree.insert(13)
        self.assertEquals(tree.search(6).value, 6)
        self.assertEquals(tree.search(2), None)
        self.assertEquals(tree.search(15), None)
        self.assertEquals(tree.search(13).value, 13)
        self.assertEquals(tree.search(1).value, 1)

    def test_deletion(self):
        tree = BinarySearchTree()
        tree.insert(8)
        tree.insert(3)
        tree.insert(10)
        tree.insert(1)
        tree.insert(6)
        tree.insert(4)
        tree.insert(7)
        tree.insert(14)
        tree.insert(13)

        tree.delete(1)
        tree.delete(14)
        tree.delete(6)

        self.assertEquals(tree.root.value, 8)
        self.assertEquals(tree.root.left.value, 3)
        self.assertEquals(tree.root.left.left, None)
        self.assertEquals(tree.root.left.right.value, 7)
        self.assertEquals(tree.root.left.right.left.value, 4)
        self.assertEquals(tree.root.left.right.right, None)
        self.assertEquals(tree.root.right.value, 10)
        self.assertEquals(tree.root.right.right.value, 13)

        tree.delete(8)
        self.assertEquals(tree.root.value, 10)
        self.assertEquals(tree.root.right.value, 13)


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