-
Notifications
You must be signed in to change notification settings - Fork 0
/
multimap.py
61 lines (55 loc) · 1.55 KB
/
multimap.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
class MultiMap:
"""
A multimap class built upon use of an underlying map for storage.
"""
# Map type; can be redefined by subclass
_MapType = dict
def __init__(self):
"""
Create a new empty multimap instance.
"""
# create map instance for storage
self._map = self._MapType()
self._n = 0
def __iter__(self):
"""
Iterate through all (k, v) pairs in multimap.
"""
for k, secondary in self._map.items():
for v in secondary:
yield (k, v)
def add(self, k, v):
"""
Add pair (k, v) to multimap.
"""
# create empty list, if needed
container = self._map.setdefault(k, [])
container.append(v)
self._n += 1
def pop(self, k):
"""
Remove and return arbitrary (k, v) with key k (or raise KeyError).
"""
# may raise KeyError
secondary = self._map[k]
v = secondary.pop()
if len(secondary) == 0:
# no pairs left
del self._map[k]
self._n -= 1
return k, v
def find(self, k):
"""
Return arbitrary (k, v) pair with given key (or raise KeyError).
"""
# may raise KeyError
secondary = self._map[k]
return k, secondary[0]
def find_all(self, k):
"""
Generate iteration of all (k, v) pairs with given key.
"""
# empty list, by default
secondary = self._map.get(k, [])
for v in secondary:
yield (k, v)