Mercurial > public > think_complexity

diff redblacktree.py @ 20:0326803882ad
Finally completing Ch. 3, exercise 5: write a TreeMap that uses a red-black tree.
author: Brian Neal <bgneal@gmail.com>
date: Thu, 27 Dec 2012 15:30:49 -0600
parents: 3c74185c5047
--- a/redblacktree.py	Thu Dec 27 13:46:12 2012 -0600
+++ b/redblacktree.py	Thu Dec 27 15:30:49 2012 -0600
@@ -370,6 +370,21 @@
 
         return remove_flag
 
+    def find(self, key):
+        """Looks up the key in the tree and returns the corresponding value, or
+        raises a KeyError if it does not exist in the tree.
+
+        """
+        p = self.root
+        while p:
+            if p.key == key:
+                return p.value
+            else:
+                d = p.key < key
+                p = p.link[d]
+
+        raise KeyError
+
 
 if __name__ == '__main__':
     import random