changes

3 files changed, 95 insertions, 33 deletions

diff --git a/collect/bidict.py b/collect/bidict.py
index e162179..1fa66dc 100644
--- a/collect/bidict.py
+++ b/collect/bidict.py
@@ -35,14 +35,16 @@ class bidict(dict):
 
     def __setitem__(self, key, value):
         if key in self:
-            self.inverse[self[key]].remove(key)
+            old_value = self[key]
+            self.inverse[old_value].remove(key)
         super().__setitem__(key, value)
         self.inverse.setdefault(value, []).append(key)
 
     def __delitem__(self, key):
-        self.inverse.setdefault(self[key], []).remove(key)
-        if self[key] in self.inverse and not self.inverse[self[key]]:
-            del self.inverse[self[key]]
+        value = self[key]
+        self.inverse.setdefault(value, []).remove(key)
+        if value in self.inverse and not self.inverse[value]:
+            del self.inverse[value]
         super().__delitem__(key)
 
 
diff --git a/collect/bst.py b/collect/bst.py
index b4d25b3..72a3b77 100644
--- a/collect/bst.py
+++ b/collect/bst.py
@@ -5,6 +5,10 @@ from typing import Any, Optional, List
 
 class Node(object):
     def __init__(self, value: Any) -> None:
+        """
+        Note: value can be anything as long as it is comparable.
+        Check out @functools.total_ordering.
+        """
         self.left = None
         self.right = None
         self.value = value
@@ -82,10 +86,8 @@ class BinarySearchTree(object):
             return node
         elif (value < node.value and node.left is not None):
             return self._find(value, node.left)
-        else:
-            assert value > node.value
-            if node.right is not None:
-                return self._find(value, node.right)
+        elif (value > node.value and node.right is not None):
+            return self._find(value, node.right)
         return None
 
     def _parent_path(self, current: Node, target: Node):
@@ -104,7 +106,10 @@ class BinarySearchTree(object):
 
     def parent_path(self, node: Node) -> Optional[List[Node]]:
         """Return a list of nodes representing the path from
-        the tree's root to the node argument.
+        the tree's root to the node argument.  If the node does
+        not exist in the tree for some reason, the last element
+        on the path will be None but the path will indicate the
+        ancestor path of that node were it inserted.
 
         >>> t = BinarySearchTree()
         >>> t.insert(50)
@@ -131,6 +136,17 @@ class BinarySearchTree(object):
         12
         4
 
+        >>> del t[4]
+        >>> for x in t.parent_path(n):
+        ...     if x is not None:
+        ...         print(x.value)
+        ...     else:
+        ...         print(x)
+        50
+        25
+        12
+        None
+
         """
         return self._parent_path(self.root, node)
 
@@ -138,15 +154,6 @@ class BinarySearchTree(object):
         """
         Delete an item from the tree and preserve the BST property.
 
-                            50
-                           /  \
-                         25    75
-                        /     /  \
-                      22    66    85
-                     /
-                   13
-
-
         >>> t = BinarySearchTree()
         >>> t.insert(50)
         >>> t.insert(75)
@@ -155,6 +162,14 @@ class BinarySearchTree(object):
         >>> t.insert(22)
         >>> t.insert(13)
         >>> t.insert(85)
+        >>> t
+        50
+        ├──25
+        │  └──22
+        │     └──13
+        └──75
+           ├──66
+           └──85
 
         >>> for value in t.iterate_inorder():
         ...     print(value)
@@ -195,6 +210,11 @@ class BinarySearchTree(object):
         50
         66
         85
+        >>> t
+        50
+        ├──25
+        └──85
+           └──66
 
         >>> t.__delitem__(99)
         False
@@ -551,7 +571,7 @@ class BinarySearchTree(object):
 
     def repr_traverse(self, padding: str, pointer: str, node: Node, has_right_sibling: bool) -> str:
         if node is not None:
-            self.viz += f'\n{padding}{pointer}{node.value}'
+            viz = f'\n{padding}{pointer}{node.value}'
             if has_right_sibling:
                 padding += "│  "
             else:
@@ -563,8 +583,10 @@ class BinarySearchTree(object):
             else:
                 pointer_left = "└──"
 
-            self.repr_traverse(padding, pointer_left, node.left, node.right is not None)
-            self.repr_traverse(padding, pointer_right, node.right, False)
+            viz += self.repr_traverse(padding, pointer_left, node.left, node.right is not None)
+            viz += self.repr_traverse(padding, pointer_right, node.right, False)
+            return viz
+        return ""
 
     def __repr__(self):
         """
@@ -590,16 +612,16 @@ class BinarySearchTree(object):
         if self.root is None:
             return ""
 
-        self.viz = f'{self.root.value}'
+        ret = f'{self.root.value}'
         pointer_right = "└──"
         if self.root.right is None:
             pointer_left = "└──"
         else:
             pointer_left = "├──"
 
-        self.repr_traverse('', pointer_left, self.root.left, self.root.left is not None)
-        self.repr_traverse('', pointer_right, self.root.right, False)
-        return self.viz
+        ret += self.repr_traverse('', pointer_left, self.root.left, self.root.left is not None)
+        ret += self.repr_traverse('', pointer_right, self.root.right, False)
+        return ret
 
 
 if __name__ == '__main__':
diff --git a/collect/trie.py b/collect/trie.py
index b9a5a1a..3e4c917 100644
--- a/collect/trie.py
+++ b/collect/trie.py
@@ -15,6 +15,7 @@ class Trie(object):
         self.root = {}
         self.end = "~END~"
         self.length = 0
+        self.viz = ''
 
     def insert(self, item: Sequence[Any]):
         """
@@ -240,7 +241,37 @@ class Trie(object):
             return None
         return [x for x in node if x != self.end]
 
-    def repr_recursive(self, node, delimiter):
+    def repr_fancy(self, padding: str, pointer: str, parent: str, node: Any, has_sibling: bool):
+        if node is None:
+            return
+        if node is not self.root:
+            ret = f'\n{padding}{pointer}'
+            if has_sibling:
+                padding += '│  '
+            else:
+                padding += '   '
+        else:
+            ret = f'{pointer}'
+
+        child_count = 0
+        for child in node:
+            if child != self.end:
+                child_count += 1
+
+        for child in node:
+            if child != self.end:
+                if child_count > 1:
+                    pointer = "├──"
+                    has_sibling = True
+                else:
+                    pointer = "└──"
+                    has_sibling = False
+                pointer += f'{child}'
+                child_count -= 1
+                ret += self.repr_fancy(padding, pointer, node, node[child], has_sibling)
+        return ret
+
+    def repr_brief(self, node, delimiter):
         """
         A friendly string representation of the contents of the Trie.
 
@@ -249,10 +280,8 @@ class Trie(object):
         >>> t.insert([10, 0, 0, 2])
         >>> t.insert([10, 10, 10, 1])
         >>> t.insert([10, 10, 10, 2])
-        >>> t.repr_recursive(t.root, '.')
+        >>> t.repr_brief(t.root, '.')
         '10.[0.0.[1, 2], 10.10.[1, 2]]'
-        >>> print(t)
-        10[00[1, 2], 1010[1, 2]]
 
         """
         child_count = 0
@@ -260,7 +289,7 @@ class Trie(object):
         for child in node:
             if child != self.end:
                 child_count += 1
-                child_rep = self.repr_recursive(node[child], delimiter)
+                child_rep = self.repr_brief(node[child], delimiter)
                 if len(child_rep) > 0:
                     my_rep += str(child) + delimiter + child_rep + ", "
                 else:
@@ -274,7 +303,7 @@ class Trie(object):
     def __repr__(self):
         """
         A friendly string representation of the contents of the Trie.  Under
-        the covers uses repr_recursive with no delimiter
+        the covers uses repr_fancy.
 
         >>> t = Trie()
         >>> t.insert([10, 0, 0, 1])
@@ -282,10 +311,19 @@ class Trie(object):
         >>> t.insert([10, 10, 10, 1])
         >>> t.insert([10, 10, 10, 2])
         >>> print(t)
-        10[00[1, 2], 1010[1, 2]]
+        *
+        └──10
+           ├──0
+           │  └──0
+           │     ├──1
+           │     └──2
+           └──10
+              └──10
+                 ├──1
+                 └──2
 
         """
-        return self.repr_recursive(self.root, '')
+        return self.repr_fancy('', '*', self.root, self.root, False)
 
 
 if __name__ == '__main__':