Mercurial > public > think_complexity
annotate ch3ex5.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 | |
| children |
| rev | line source |
|---|---|
| bgneal@20 | 1 """Chapter 3, exercise 5. |
| bgneal@20 | 2 |
| bgneal@20 | 3 "A drawback of hashtables is that the elements have to be hashable, which usually |
| bgneal@20 | 4 means they have to be immutable. That's why, in Python, you can use tuples but |
| bgneal@20 | 5 not lists as keys in a dictionary. An alternative is to use a tree-based map. |
| bgneal@20 | 6 Write an implementation of the map interface called TreeMap that uses |
| bgneal@20 | 7 a red-black tree to perform add and get in log time." |
| bgneal@20 | 8 |
| bgneal@20 | 9 |
| bgneal@20 | 10 """ |
| bgneal@20 | 11 import redblacktree |
| bgneal@20 | 12 |
| bgneal@20 | 13 |
| bgneal@20 | 14 class TreeMap(object): |
| bgneal@20 | 15 """A tree-based map class.""" |
| bgneal@20 | 16 |
| bgneal@20 | 17 def __init__(self): |
| bgneal@20 | 18 self.tree = redblacktree.Tree() |
| bgneal@20 | 19 |
| bgneal@20 | 20 def get(self, k): |
| bgneal@20 | 21 """Looks up the key (k) and returns the corresponding value, or raises |
| bgneal@20 | 22 a KeyError if the key is not found. |
| bgneal@20 | 23 |
| bgneal@20 | 24 """ |
| bgneal@20 | 25 return self.tree.find(k) |
| bgneal@20 | 26 |
| bgneal@20 | 27 def add(self, k, v): |
| bgneal@20 | 28 """Adds the key/value pair (k, v) to the tree. If the key already |
| bgneal@20 | 29 exists, the value is updated to the new value v. |
| bgneal@20 | 30 |
| bgneal@20 | 31 Returns True if the pair was inserted, and False if the key already |
| bgneal@20 | 32 existed and the tree was updated. |
| bgneal@20 | 33 |
| bgneal@20 | 34 """ |
| bgneal@20 | 35 node, inserted = self.tree.insert(k, v) |
| bgneal@20 | 36 if not inserted: |
| bgneal@20 | 37 node.value = v |
| bgneal@20 | 38 return inserted |
| bgneal@20 | 39 |
| bgneal@20 | 40 def remove(self, k): |
| bgneal@20 | 41 """Removes the mapping with the given key (k). |
| bgneal@20 | 42 Raises a KeyError if the mapping was not found. |
| bgneal@20 | 43 |
| bgneal@20 | 44 """ |
| bgneal@20 | 45 result = self.tree.remove(k) |
| bgneal@20 | 46 if not result: |
| bgneal@20 | 47 raise KeyError |
| bgneal@20 | 48 |
| bgneal@20 | 49 def __len__(self): |
| bgneal@20 | 50 """Returns the number of mappings in the map.""" |
| bgneal@20 | 51 return len(self.tree) |
| bgneal@20 | 52 |
| bgneal@20 | 53 |
| bgneal@20 | 54 def main(script): |
| bgneal@20 | 55 import string |
| bgneal@20 | 56 m = TreeMap() |
| bgneal@20 | 57 s = string.ascii_lowercase |
| bgneal@20 | 58 |
| bgneal@20 | 59 for k, v in enumerate(s): |
| bgneal@20 | 60 m.add([k], v) |
| bgneal@20 | 61 |
| bgneal@20 | 62 for k in range(len(s)): |
| bgneal@20 | 63 key = [k] |
| bgneal@20 | 64 print key, m.get(key) |
| bgneal@20 | 65 m.remove(key) |
| bgneal@20 | 66 |
| bgneal@20 | 67 assert len(m) == 0 |
| bgneal@20 | 68 |
| bgneal@20 | 69 if __name__ == '__main__': |
| bgneal@20 | 70 import sys |
| bgneal@20 | 71 main(*sys.argv) |