Skip to content
Newer
Older
100644 156 lines (131 sloc) 4.41 KB
d4ac95f @dhain initial commit
authored Oct 23, 2009
1 class NeedMore(Exception):
2 pass
3
4
5 class Node(object):
6f8aca6 @dhain add some docstrings
authored Oct 25, 2009
6 """Internal representation of Trie nodes."""
d4ac95f @dhain initial commit
authored Oct 23, 2009
7 __slots__ = 'parent key nodes value'.split()
8 no_value = object()
9
10 def __init__(self, parent, key, nodes, value):
11 self.parent = parent
12 self.key = key
13 self.nodes = nodes
14 self.value = value
15
16 @property
17 def keypath(self):
18 n = self
19 keypath = [n.key for n in iter(lambda: n.parent, None) if n.key]
20 keypath.reverse()
21 keypath.append(self.key)
22 return keypath
23
24 def walk(self):
25 nodes = [self]
26 while nodes:
27 node = nodes.pop()
28 if node.value is not node.no_value:
29 yield node
30 nodes.extend(node.nodes[key] for key in sorted(node.nodes, reverse=True))
31
32
33 class Trie(object):
6f8aca6 @dhain add some docstrings
authored Oct 25, 2009
34 """A simple prefix tree (trie) implementation.
35
36 If attempting to access a node without a value, but with descendents,
37 NeedMore will be raised. If there are no descendents, KeyError will be
38 raised.
39
40 Usage:
41
42 >>> import trie
43 >>> from pprint import pprint
44 >>> t = trie.Trie()
45 >>> t['foobaz'] = 'Here is a foobaz.'
46 >>> t['foobar'] = 'This is a foobar.'
47 >>> t['fooqat'] = "What's a fooqat?"
48 >>> pprint(list(t))
49 [['f', 'o', 'o', 'b', 'a', 'r'],
50 ['f', 'o', 'o', 'b', 'a', 'z'],
51 ['f', 'o', 'o', 'q', 'a', 't']]
52 >>> pprint(list(t.iteritems()))
53 [(['f', 'o', 'o', 'b', 'a', 'r'], 'This is a foobar.'),
54 (['f', 'o', 'o', 'b', 'a', 'z'], 'Here is a foobaz.'),
55 (['f', 'o', 'o', 'q', 'a', 't'], "What's a fooqat?")]
56 >>> t['foo']
57 Traceback (most recent call last):
58 ...
59 NeedMore
60 >>> t['fooqux']
61 Traceback (most recent call last):
62 ...
63 KeyError: 'fooqux'
64 >>> t.children('fooba')
65 {'r': 'This is a foobar.', 'z': 'Here is a foobaz.'}
66 >>> del t['foobaz']
67 >>> pprint(list(t.iteritems()))
68 [(['f', 'o', 'o', 'b', 'a', 'r'], 'This is a foobar.'),
69 (['f', 'o', 'o', 'q', 'a', 't'], "What's a fooqat?")]
70 """
71
d4ac95f @dhain initial commit
authored Oct 23, 2009
72 def __init__(self, root_data=Node.no_value, mapping=()):
6f8aca6 @dhain add some docstrings
authored Oct 25, 2009
73 """Initialize a Trie instance.
74
75 Args (both optional):
76 root_data: value of the root node (ie. Trie('hello')[()] == 'hello').
77 mapping: a sequence of (key, value) pairs to initialize with.
78 """
d4ac95f @dhain initial commit
authored Oct 23, 2009
79 self.root = Node(None, None, {}, root_data)
80 self.extend(mapping)
81
82 def extend(self, mapping):
6f8aca6 @dhain add some docstrings
authored Oct 25, 2009
83 """Update the Trie with a sequence of (key, value) pairs."""
d4ac95f @dhain initial commit
authored Oct 23, 2009
84 for k, v in mapping:
85 self[k] = v
86
87 def __setitem__(self, k, v):
88 n = self.root
89 for c in k:
90 n = n.nodes.setdefault(c, Node(n, c, {}, Node.no_value))
91 n.value = v
92
93 def _getnode(self, k):
94 n = self.root
95 for c in k:
96 try:
97 n = n.nodes[c]
98 except KeyError:
99 raise KeyError(k)
100 return n
101
102 def __getitem__(self, k):
103 n = self._getnode(k)
104 if n.value is Node.no_value:
105 if n.nodes:
106 raise NeedMore()
107 else:
108 raise KeyError(k)
109 return n.value
110
111 def __delitem__(self, k):
112 n = self._getnode(k)
113 if n.value is Node.no_value:
114 raise KeyError(k)
115 n.value = Node.no_value
116 while True:
fc360b5 fix bug in deleting notes
Wil Mahan authored Aug 9, 2013
117 if n.nodes or not n.parent or n.value is not Node.no_value:
d4ac95f @dhain initial commit
authored Oct 23, 2009
118 break
119 del n.parent.nodes[n.key]
120 n = n.parent
121
122 def children(self, k):
6f8aca6 @dhain add some docstrings
authored Oct 25, 2009
123 """Return a dict of the immediate children of the given key.
124
125 Example:
126 >>> t = Trie()
127 >>> t['foobaz'] = 'Here is a foobaz.'
128 >>> t['foobar'] = 'This is a foobar.'
129 >>> t.children('fooba')
130 {'r': 'This is a foobar.', 'z': 'Here is a foobaz.'}
131 """
d4ac95f @dhain initial commit
authored Oct 23, 2009
132 n = self._getnode(k)
133 return dict((k, n.nodes[k].value)
134 for k in n.nodes
135 if n.nodes[k].value is not Node.no_value)
136
137 def __iter__(self):
6f8aca6 @dhain add some docstrings
authored Oct 25, 2009
138 """Yield the keys in order."""
d4ac95f @dhain initial commit
authored Oct 23, 2009
139 for node in self.root.walk():
140 yield node.keypath
141
142 def iteritems(self):
6f8aca6 @dhain add some docstrings
authored Oct 25, 2009
143 """Yield (key, value) pairs in order."""
d4ac95f @dhain initial commit
authored Oct 23, 2009
144 for node in self.root.walk():
145 yield node.keypath, node.value
146
147 def itervalues(self):
6f8aca6 @dhain add some docstrings
authored Oct 25, 2009
148 """Yield values in order."""
d4ac95f @dhain initial commit
authored Oct 23, 2009
149 for node in self.root.walk():
150 yield node.value
6f8aca6 @dhain add some docstrings
authored Oct 25, 2009
151
152
153 if __name__ == '__main__':
154 import doctest
155 doctest.testmod()
Something went wrong with that request. Please try again.