/
_pbag.py
267 lines (212 loc) · 6.6 KB
/
_pbag.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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
from collections import Container, Iterable, Sized, Hashable
from functools import reduce
from pyrsistent._pmap import pmap
def _add_to_counters(counters, element):
return counters.set(element, counters.get(element, 0) + 1)
class PBag(object):
"""
A persistent bag/multiset type.
Requires elements to be hashable, and allows duplicates, but has no
ordering. Bags are hashable.
Do not instantiate directly, instead use the factory functions :py:func:`b`
or :py:func:`pbag` to create an instance.
Some examples:
>>> s = pbag([1, 2, 3, 1])
>>> s2 = s.add(4)
>>> s3 = s2.remove(1)
>>> s
pbag([1, 1, 2, 3])
>>> s2
pbag([1, 1, 2, 3, 4])
>>> s3
pbag([1, 2, 3, 4])
"""
__slots__ = ('_counts', '__weakref__')
def __init__(self, counts):
self._counts = counts
def add(self, element):
"""
Add an element to the bag.
>>> s = pbag([1])
>>> s2 = s.add(1)
>>> s3 = s.add(2)
>>> s2
pbag([1, 1])
>>> s3
pbag([1, 2])
"""
return PBag(_add_to_counters(self._counts, element))
def update(self, iterable):
"""
Update bag with all elements in iterable.
>>> s = pbag([1])
>>> s.update([1, 2])
pbag([1, 1, 2])
"""
if iterable:
return PBag(reduce(_add_to_counters, iterable, self._counts))
return self
def remove(self, element):
"""
Remove an element from the bag.
>>> s = pbag([1, 1, 2])
>>> s2 = s.remove(1)
>>> s3 = s.remove(2)
>>> s2
pbag([1, 2])
>>> s3
pbag([1, 1])
"""
if element not in self._counts:
raise KeyError(element)
elif self._counts[element] == 1:
newc = self._counts.remove(element)
else:
newc = self._counts.set(element, self._counts[element] - 1)
return PBag(newc)
def count(self, element):
"""
Return the number of times an element appears.
>>> pbag([]).count('non-existent')
0
>>> pbag([1, 1, 2]).count(1)
2
"""
return self._counts.get(element, 0)
def __len__(self):
"""
Return the length including duplicates.
>>> len(pbag([1, 1, 2]))
3
"""
return sum(self._counts.itervalues())
def __iter__(self):
"""
Return an iterator of all elements, including duplicates.
>>> list(pbag([1, 1, 2]))
[1, 1, 2]
>>> list(pbag([1, 2]))
[1, 2]
"""
for elt, count in self._counts.iteritems():
for i in range(count):
yield elt
def __contains__(self, elt):
"""
Check if an element is in the bag.
>>> 1 in pbag([1, 1, 2])
True
>>> 0 in pbag([1, 2])
False
"""
return elt in self._counts
def __repr__(self):
return "pbag({0})".format(list(self))
def __eq__(self, other):
"""
Check if two bags are equivalent, honoring the number of duplicates,
and ignoring insertion order.
>>> pbag([1, 1, 2]) == pbag([1, 2])
False
>>> pbag([2, 1, 0]) == pbag([0, 1, 2])
True
"""
if type(other) is not PBag:
raise TypeError("Can only compare PBag with PBags")
return self._counts == other._counts
def __lt__(self, other):
raise TypeError('PBags are not orderable')
__le__ = __lt__
__gt__ = __lt__
__ge__ = __lt__
# Multiset-style operations similar to collections.Counter
def __add__(self, other):
"""
Combine elements from two PBags.
>>> pbag([1, 2, 2]) + pbag([2, 3, 3])
pbag([1, 2, 2, 2, 3, 3])
"""
if not isinstance(other, PBag):
return NotImplemented
result = self._counts.evolver()
for elem, other_count in other._counts.iteritems():
result[elem] = self.count(elem) + other_count
return PBag(result.persistent())
def __sub__(self, other):
"""
Remove elements from one PBag that are present in another.
>>> pbag([1, 2, 2, 2, 3]) - pbag([2, 3, 3, 4])
pbag([1, 2, 2])
"""
if not isinstance(other, PBag):
return NotImplemented
result = self._counts.evolver()
for elem, other_count in other._counts.iteritems():
newcount = self.count(elem) - other_count
if newcount > 0:
result[elem] = newcount
elif elem in self:
result.remove(elem)
return PBag(result.persistent())
def __or__(self, other):
"""
Union: Keep elements that are present in either of two PBags.
>>> pbag([1, 2, 2, 2]) | pbag([2, 3, 3])
pbag([1, 2, 2, 2, 3, 3])
"""
if not isinstance(other, PBag):
return NotImplemented
result = self._counts.evolver()
for elem, other_count in other._counts.iteritems():
count = self.count(elem)
newcount = max(count, other_count)
result[elem] = newcount
return PBag(result.persistent())
def __and__(self, other):
"""
Intersection: Only keep elements that are present in both PBags.
>>> pbag([1, 2, 2, 2]) & pbag([2, 3, 3])
pbag([2])
"""
if not isinstance(other, PBag):
return NotImplemented
result = pmap().evolver()
for elem, count in self._counts.iteritems():
newcount = min(count, other.count(elem))
if newcount > 0:
result[elem] = newcount
return PBag(result.persistent())
def __hash__(self):
"""
Hash based on value of elements.
>>> m = pmap({pbag([1, 2]): "it's here!"})
>>> m[pbag([2, 1])]
"it's here!"
>>> pbag([1, 1, 2]) in m
False
"""
return hash(self._counts)
Container.register(PBag)
Iterable.register(PBag)
Sized.register(PBag)
Hashable.register(PBag)
def b(*elements):
"""
Construct a persistent bag.
Takes an arbitrary number of arguments to insert into the new persistent
bag.
>>> b(1, 2, 3, 2)
pbag([1, 2, 2, 3])
"""
return pbag(elements)
def pbag(elements):
"""
Convert an iterable to a persistent bag.
Takes an iterable with elements to insert.
>>> pbag([1, 2, 3, 2])
pbag([1, 2, 2, 3])
"""
if not elements:
return _EMPTY_PBAG
return PBag(reduce(_add_to_counters, elements, pmap()))
_EMPTY_PBAG = PBag(pmap())