bgneal@20: """Chapter 3, exercise 5.
bgneal@20: 
bgneal@20: "A drawback of hashtables is that the elements have to be hashable, which usually
bgneal@20: means they have to be immutable. That's why, in Python, you can use tuples but
bgneal@20: not lists as keys in a dictionary. An alternative is to use a tree-based map.
bgneal@20: Write an implementation of the map interface called TreeMap that uses
bgneal@20: a red-black tree to perform add and get in log time."
bgneal@20: 
bgneal@20: 
bgneal@20: """
bgneal@20: import redblacktree
bgneal@20: 
bgneal@20: 
bgneal@20: class TreeMap(object):
bgneal@20:     """A tree-based map class."""
bgneal@20: 
bgneal@20:     def __init__(self):
bgneal@20:         self.tree = redblacktree.Tree()
bgneal@20: 
bgneal@20:     def get(self, k):
bgneal@20:         """Looks up the key (k) and returns the corresponding value, or raises
bgneal@20:         a KeyError if the key is not found.
bgneal@20: 
bgneal@20:         """
bgneal@20:         return self.tree.find(k)
bgneal@20: 
bgneal@20:     def add(self, k, v):
bgneal@20:         """Adds the key/value pair (k, v) to the tree. If the key already
bgneal@20:         exists, the value is updated to the new value v.
bgneal@20: 
bgneal@20:         Returns True if the pair was inserted, and False if the key already
bgneal@20:         existed and the tree was updated.
bgneal@20: 
bgneal@20:         """
bgneal@20:         node, inserted = self.tree.insert(k, v)
bgneal@20:         if not inserted:
bgneal@20:             node.value = v
bgneal@20:         return inserted
bgneal@20: 
bgneal@20:     def remove(self, k):
bgneal@20:         """Removes the mapping with the given key (k).
bgneal@20:         Raises a KeyError if the mapping was not found.
bgneal@20: 
bgneal@20:         """
bgneal@20:         result = self.tree.remove(k)
bgneal@20:         if not result:
bgneal@20:             raise KeyError
bgneal@20: 
bgneal@20:     def __len__(self):
bgneal@20:         """Returns the number of mappings in the map."""
bgneal@20:         return len(self.tree)
bgneal@20: 
bgneal@20: 
bgneal@20: def main(script):
bgneal@20:     import string
bgneal@20:     m = TreeMap()
bgneal@20:     s = string.ascii_lowercase
bgneal@20: 
bgneal@20:     for k, v in enumerate(s):
bgneal@20:         m.add([k], v)
bgneal@20: 
bgneal@20:     for k in range(len(s)):
bgneal@20:         key = [k]
bgneal@20:         print key, m.get(key)
bgneal@20:         m.remove(key)
bgneal@20: 
bgneal@20:     assert len(m) == 0
bgneal@20: 
bgneal@20: if __name__ == '__main__':
bgneal@20:     import sys
bgneal@20:     main(*sys.argv)