Move height and insert into Node. Other small fixes.

leftrb/bst.py

@@ -51,6 +51,24 @@ def search(self, key):
                 return self.right.search(key)
             return None
 
+        def insert(self, key, value=None):
+            """
+            Insert a node recursively.
+            """
+            if self.key == key:
+                self.val = value
+            elif key < self.key:
+                if self.left is None:
+                    self.left = self.__class__(key, value)
+                else:
+                    self.left = self.left.insert(key, value)
+            else:
+                if self.right is None:
+                    self.right = self.__class__(key, value)
+                else:
+                    self.right = self.right.insert(key, value)
+            return self
+
         def min(self):
             """
             Smallest node in the subtree.
@@ -67,27 +85,11 @@ def search(self, key):
         """
         Search the tree with a key. Return a value or None.
         """
-        return self.root.search(key)
+        return self.root.search(key) if self.root is not None else None
 
     def insert(self, key, value=None):
         """
         Insert a key with optional value into tree.
         """
-        self.root = self._insert(self.root, key, value)
-
-    @classmethod
-    def _insert(cls, h, key, value=None):
-        """
-        Insert a node (h) node recursively.
-        """
-        if h is None:
-            return cls.Node(key, value)
-
-        if h.key == key:
-            h.val = value
-        elif key < h.key:
-            h.left = cls._insert(h.left, key, value)
-        else:
-            h.right = cls._insert(h.right, key, value)
+        self.root = self.Node(key, value) if self.root is None else self.root.insert(key, value)
 
-        return h

leftrb/llrb.py

@@ -30,7 +30,7 @@
 """
 
 import sys
-from bst import BinarySearchTree
+from leftrb.bst import BinarySearchTree
 
 
 __all__ = ['LeftRB']
@@ -73,6 +73,23 @@ def __init__(self, key, val=None):
             self.color = RED  # new nodes are always red
             self.height = 1
 
+        def insert(self, key, value=None):
+            """
+            Recursively insert a node with key and optional value into the tree below.
+            """
+            # Move this to the end to get 2-3 trees
+            if is_red(self.left) and is_red(self.right):
+                LeftRB._flip_colors(self)
+
+            self = super(LeftRB.Node, self).insert(key, value)
+
+            if is_red(self.right) and is_black(self.left):
+                self = LeftRB._rotate_left(self)
+            if is_red(self.left) and self.left and is_red(self.left.left):
+                self = LeftRB._rotate_right(self)
+
+            return LeftRB._setHeight(self)
+
         def size(self):
             """
             Number of nodes in the subtree below node.
@@ -81,13 +98,12 @@ def size(self):
             return 1 + sum(map(lambda child: child.size(), filter(None, [self.left, self.right])))
 
         def __repr__(self):
-            return "<{0} at {1}, key={2}, value={3}, color={4}, N={5}, height={6}>".format(
+            return "<{0} at {1}, key={2}, value={3}, color={4}, height={5}>".format(
                 self.__class__.__name__,
                 id(self),
                 self.key,
                 self.val,
                 'red' if is_red(self) else 'black',
-                self.N,
                 self.height
             )
 
@@ -107,20 +123,13 @@ def __len__(self):
         """
         Number of nodes in the tree.
         """
-        return 0 if not self.root else self.root.size()
+        return 0 if self.root is None else self.root.size()
 
     def height(self):
         """
         Height of the tree.
         """
-        return self._height(self.root)
-
-    @staticmethod
-    def _height(x):
-        """
-        The height of the tree below node (x).
-        """
-        return 0 if x is None else x.height
+        return 0 if self.root is None else self.root.height
 
     def min(self):
         """
@@ -138,31 +147,9 @@ def insert(self, key, value=None):
         """
         Insert a key with optional value into the tree.
         """
-        self.root = self._insert(self.root, key, value)
+        super(LeftRB, self).insert(key, value)
         self.root.color = BLACK
 
-    @classmethod
-    def _insert(cls, h, key, value=None):
-        """
-        Recursively insert a node with key and optional value
-        into the tree below node (h).
-        """
-        if h is None:
-            return cls.Node(key, value)
-
-        # Move this to the end to get 2-3 trees
-        if is_red(h.left) and is_red(h.right):
-            cls._flip_colors(h)
-
-        h = super(cls, cls)._insert(h, key, value)
-
-        if is_red(h.right) and is_black(h.left):
-            h = cls._rotate_left(h)
-        if is_red(h.left) and h.left and is_red(h.left.left):
-            h = cls._rotate_right(h)
-
-        return cls._setHeight(h)
-
     def delete(self, key):
         """
         Delete a node with the given key from the tree.
@@ -200,8 +187,8 @@ def _delete(self, h, key):
                 h = self._move_red_right(h)
 
             if key == h.key:
-                h.value = self._search(h.right, self._min(h.right))
-                h.key = self._min(h.right)
+                h.value = h.right.search(h.right.min())
+                h.key = h.right.min()
                 h.right = self._delete_min(h.right)
             else:
                 h.right = self._delete(h.right, key)
@@ -351,7 +338,7 @@ def _setHeight(cls, h):
         """
         Update size and height of node (h).
         """
-        h.height = max(cls._height(h.left), cls._height(h.right)) + 1
+        h.height = max(h.left and h.left.height or 0, h.right and h.right.height or 0) + 1
         return h
 
 

leftrb/test/test_bst.py

@@ -16,7 +16,7 @@ class TestBST(object):
 
     def test_insert_and_search(self):
         t = Tree()
-        r = random.sample(xrange(1000), 900)
+        r = random.sample(range(1000), 900)
         for x in r:
             t.insert(x)
         res = [t.search(x) for x in r]

leftrb/test/test_llrb.py

@@ -4,6 +4,7 @@
 from __future__ import absolute_import
 
 import pytest
+import math
 import random
 
 from leftrb.test.test_bst import TestBST as Base
@@ -21,7 +22,7 @@ def fill_tree(items):
 
 class TestLeftRB(Base):
 
-    items = [5, 1, 3, 6, 2, 7]
+    items = [5, 1, 3, 6]
 
     def test_search(self):
         t = fill_tree(self.items)
@@ -37,6 +38,13 @@ def test_len(self):
         t = fill_tree(self.items)
         assert len(t) == len(self.items)
 
+    def test_height(self):
+        for n in range(16):
+            items = random.sample(range(n), n)
+            t = fill_tree(items)
+            print("Items: {0}".format(items))
+            assert t.height() <= int(math.ceil(2 * math.log(n + 1, 2)))
+
     def test_min(self):
         t = fill_tree(self.items)
         assert t.min() == min(self.items)
@@ -46,7 +54,7 @@ def test_max(self):
         assert t.max() == max(self.items)
 
     def test_delete(self):
-        t = fill_tree(random.sample(xrange(100), 90))
+        t = fill_tree(random.sample(range(100), 90))
 
         key = random.randint(0, 999)
         value = str(key)