"""Bag class definitions."""
import heapq
from abc import ABCMeta, abstractmethod
from collections import defaultdict
from collections.abc import Collection, Hashable, Set
from operator import itemgetter
from ._util import deprecated
__all__ = (
'Bag',
'BagView',
'CountsView',
'UniqueElementsView',
'bag',
'frozenbag',
)
class BagView(Collection):
"""Base class for bag views."""
__metaclass__ = ABCMeta
__slots__ = ('bag', )
def __init__(self, bag):
self.bag = bag
def __repr__(self):
return '{0.__class__.__name__}({0.bag!r})'.format(self)
def __len__(self):
return self.bag.num_unique_elements()
@abstractmethod
def __iter__(self):
raise NotImplementedError
@abstractmethod
def __contains__(self, elem):
raise NotImplementedError
[docs]class UniqueElementsView(BagView):
"""A view for the unique items and their counts in a bag.
.. versionadded:: 1.0
"""
def __iter__(self):
for elem in self.bag._dict:
yield elem
def __contains__(self, elem):
return elem in self.bag
[docs]class CountsView(BagView):
"""A view for the unique items and their counts in a bag.
.. versionadded:: 1.0
"""
__slots__ = ('bag', )
def __len__(self):
return self.bag.num_unique_elements()
def __iter__(self):
for elem in self.bag.unique_elements():
yield elem, self.bag.count(elem)
def __contains__(self, item):
elem, count = item
return self.bag.count(elem) == count
[docs]class Bag(Collection):
"""Base class for bag classes.
Base class for bag and frozenbag. Is not mutable and not hashable, so there's
no reason to use this instead of either bag or frozenbag.
"""
# Basic object methods
[docs] def __init__(self, iterable=None):
"""Create a new bag.
If iterable isn't given, is None or is empty then the bag starts empty.
Otherwise each element from iterable will be added to the bag
however many times it appears.
This runs in O(len(iterable))
"""
self._dict = dict()
self._size = 0
if iterable:
if isinstance(iterable, Bag):
self._dict = iterable._dict.copy()
self._size = iterable._size
else:
for value in iterable:
self._increment_count(value)
def _set_count(self, elem, count):
if count < 0:
raise ValueError
self._size += count - self.count(elem)
if count == 0:
self._dict.pop(elem, None)
else:
self._dict[elem] = count
def _increment_count(self, elem, count=1):
self._set_count(elem, self.count(elem) + count)
@classmethod
def _from_iterable(cls, it):
return cls(it)
[docs] def copy(self):
"""Create a shallow copy of self.
This runs in O(len(self.num_unique_elements()))
"""
out = self._from_iterable(None)
out._dict = self._dict.copy()
out._size = self._size
return out
def __repr__(self):
if self._size == 0:
return '{0}()'.format(self.__class__.__name__)
else:
repr_format = '{class_name}({values!r})'
return repr_format.format(
class_name=self.__class__.__name__,
values=tuple(self),
)
def __str__(self):
if self._size == 0:
return '{class_name}()'.format(class_name=self.__class__.__name__)
else:
format_single = '{elem!r}'
format_mult = '{elem!r}^{mult}'
strings = []
for elem, mult in self.counts():
if mult > 1:
strings.append(format_mult.format(elem=elem, mult=mult))
else:
strings.append(format_single.format(elem=elem))
return '{%s}' % ', '.join(strings)
# New public methods (not overriding/implementing anything)
[docs] def num_unique_elements(self):
"""Return the number of unique elements.
This runs in O(1) time
"""
return len(self._dict)
[docs] def unique_elements(self):
"""Return a view of unique elements in this bag."""
return UniqueElementsView(self)
[docs] def count(self, value):
"""Return the number of value present in this bag.
If value is not in the bag no Error is raised, instead 0 is returned.
This runs in O(1) time
Args:
value: The element of self to get the count of
Returns:
int: The count of value in self
"""
return self._dict.get(value, 0)
[docs] @deprecated(
"Use `heapq.nlargest(n, self.counts(), key=itemgetter(1))` instead or "
"`sorted(self.counts(), reverse=True, key=itemgetter(1))` for `n=None`",
'1.0',
)
def nlargest(self, n=None):
"""List the n most common elements and their counts.
List is from the most
common to the least. If n is None, the list all element counts.
Run time should be O(m log m) where m is len(self)
Args:
n (int): The number of elements to return
"""
if n is None:
return sorted(self.counts(), key=itemgetter(1), reverse=True)
else:
return heapq.nlargest(n, self.counts(), key=itemgetter(1))
[docs] def counts(self):
"""Return a view of the unique elements in self and their counts.
.. versionadded:: 1.0.3
"""
return CountsView(self)
[docs] @classmethod
def from_mapping(cls, mapping):
"""Create a bag from a dict of elem->count.
Each key in the dict is added if the value is > 0.
Raises:
ValueError: If any count is < 0.
"""
out = cls()
for elem, count in mapping.items():
out._set_count(elem, count)
return out
# implementing Sized methods
[docs] def __len__(self):
"""Return the cardinality of the bag.
This runs in O(1)
"""
return self._size
# implementing Container methods
[docs] def __contains__(self, value):
"""Return the multiplicity of the element.
This runs in O(1)
"""
return self.count(value)
# implementing Iterable methods
[docs] def __iter__(self):
"""Iterate through all elements.
Multiple copies will be returned if they exist.
"""
for value, count in self.counts():
for _ in range(count):
yield value
# Comparison methods
[docs] def issubset(self, other):
"""Check that every element in self has a count <= in other.
Args:
other (Iterable)
"""
if not isinstance(other, Bag):
return self.issubset(frozenbag(other))
for elem, count in self.counts():
if not count <= other.count(elem):
return False
return True
[docs] def issuperset(self, other):
"""Check that every element in self has a count >= in other.
Args:
other (Iterable)
"""
if not isinstance(other, Bag):
return self.issuperset(bag(other))
for elem, count in other.counts():
if not self.count(elem) >= count:
return False
return True
def __le__(self, other):
if not isinstance(other, Bag):
return NotImplemented
return len(self) <= len(other) and self.issubset(other)
def __lt__(self, other):
if not isinstance(other, Bag):
return NotImplemented
return len(self) < len(other) and self.issubset(other)
def __gt__(self, other):
if not isinstance(other, Bag):
return NotImplemented
return len(self) > len(other) and self.issuperset(other)
def __ge__(self, other):
if not isinstance(other, Bag):
return NotImplemented
return len(self) >= len(other) and self.issuperset(other)
def __eq__(self, other):
if not isinstance(other, Bag):
return False
return self._dict == other._dict
def __ne__(self, other):
return not (self == other)
# Operations - &, |, +, -, ^, * and isdisjoint
def _iadd(self, other):
"""Add all of the elements of other to self.
if isinstance(it, Bag):
This runs in O(it.num_unique_elements())
else:
This runs in O(len(it))
"""
if isinstance(other, Bag):
for elem, count in other.counts():
self._increment_count(elem, count)
else:
for elem in other:
self._increment_count(elem, 1)
return self
def _iand(self, other):
"""Set multiplicity of each element to the minimum of the two collections.
if isinstance(other, Bag):
This runs in O(other.num_unique_elements())
else:
This runs in O(len(other))
"""
# TODO do we have to create a bag from the other first?
if not isinstance(other, Bag):
other = self._from_iterable(other)
for elem, old_count in set(self.counts()):
other_count = other.count(elem)
new_count = min(other_count, old_count)
self._set_count(elem, new_count)
return self
def _ior(self, other):
"""Set multiplicity of each element to the maximum of the two collections.
if isinstance(other, Bag):
This runs in O(other.num_unique_elements())
else:
This runs in O(len(other))
"""
# TODO do we have to create a bag from the other first?
if not isinstance(other, Bag):
other = self._from_iterable(other)
for elem, other_count in other.counts():
old_count = self.count(elem)
new_count = max(other_count, old_count)
self._set_count(elem, new_count)
return self
def _ixor(self, other):
"""Set self to the symmetric difference between the sets.
if isinstance(other, Bag):
This runs in O(other.num_unique_elements())
else:
This runs in O(len(other))
"""
if isinstance(other, Bag):
for elem, other_count in other.counts():
count = abs(self.count(elem) - other_count)
self._set_count(elem, count)
else:
# Let a = self.count(elem) and b = other.count(elem)
# if a >= b then elem is removed from self b times leaving a - b
# if a < b then elem is removed from self a times then added (b - a)
# times leaving a - a + (b - a) = b - a
for elem in other:
try:
self._increment_count(elem, -1)
except ValueError:
self._increment_count(elem, 1)
return self
def _isub(self, other):
"""Discard the elements of other from self.
if isinstance(it, Bag):
This runs in O(it.num_unique_elements())
else:
This runs in O(len(it))
"""
if isinstance(other, Bag):
for elem, other_count in other.counts():
try:
self._increment_count(elem, -other_count)
except ValueError:
self._set_count(elem, 0)
else:
for elem in other:
try:
self._increment_count(elem, -1)
except ValueError:
pass
return self
[docs] def __and__(self, other):
"""Intersection is the minimum of corresponding counts.
This runs in O(l + n) where:
* n is self.num_unique_elements()
* `l = 1` if other is a bag else `l = len(other)`
"""
return self.copy()._iand(other)
[docs] def isdisjoint(self, other):
"""Return if this bag is disjoint with the passed collection.
This runs in O(len(other))
"""
for value in other:
if value in self:
return False
return True
[docs] def __or__(self, other):
"""Union is the maximum of all elements.
This runs in O(m + n) where:
* `n = self.num_unique_elements()`
* m = other.num_unique_elements() if other is a bag else m = len(other)
"""
return self.copy()._ior(other)
[docs] def __add__(self, other):
"""Return a new bag also containing all the elements of other.
self + other = self & other + self | other
This runs in O(m + n) where:
* n is self.num_unique_elements()
* m is len(other)
Args:
other (Iterable): elements to add to self
"""
return self.copy()._iadd(other)
[docs] def __sub__(self, other):
"""Difference between the sets.
For normal sets this is all x s.t. x in self and x not in other.
For bags this is count(x) = max(0, self.count(x)-other.count(x))
This runs in O(m + n) where:
* n is self.num_unique_elements()
* m is len(other)
Args:
other (Iterable): elements to remove
"""
return self.copy()._isub(other)
[docs] def __mul__(self, other):
"""Cartesian product with other."""
return self.product(other)
[docs] def product(self, other, operator=None):
"""Cartesian product of the two sets.
Optionally, pass an operator to combine elements instead of creating a
tuple.
This should run in O(m*n+l) where:
* `m` is the number of unique elements in `self`
* `n` is the number of unique elements in `other`
* `l` is 0 if `other` is a bag, else `l` is the `len(other)`
Args:
other (Iterable): The iterable to take the product with.
operator (Callable): A function that accepts an element from self
and other and returns a combined value to include in the return
value.
"""
if not isinstance(other, Bag):
other = self._from_iterable(other)
values = defaultdict(int)
for elem, count in self.counts():
for other_elem, other_count in other.counts():
if operator:
new_elem = operator(elem, other_elem)
else:
new_elem = (elem, other_elem)
new_count = count * other_count
values[new_elem] += new_count
return self.from_mapping(values)
[docs] def __xor__(self, other):
"""Symmetric difference between the sets.
other can be any iterable.
This runs in O(m + n) where:
m = len(self)
n = len(other)
"""
return self.copy()._ixor(other)
[docs]class bag(Bag):
"""bag is a mutable unhashable bag.
.. automethod:: __init__
"""
[docs] def pop(self):
"""Remove and return an element of self."""
# TODO can this be done more efficiently (no need to create an iterator)?
it = iter(self)
try:
value = next(it)
except StopIteration:
raise KeyError('pop from an empty bag')
self.remove(value)
return value
[docs] def add(self, elem):
"""Add elem to self."""
self._increment_count(elem)
[docs] def discard(self, elem):
"""Remove elem from this bag, silent if it isn't present."""
try:
self.remove(elem)
except ValueError:
pass
[docs] def remove(self, elem):
"""Remove elem from this bag, raising a ValueError if it isn't present.
Args:
elem: object to remove from self
Raises:
ValueError: if the elem isn't present
"""
self._increment_count(elem, -1)
[docs] def discard_all(self, other):
"""Discard all of the elems from other."""
self._isub(other)
[docs] def remove_all(self, other):
"""Remove all of the elems from other.
Raises a ValueError if the multiplicity of any elem in other is greater
than in self.
"""
if not self.issuperset(other):
raise ValueError('Passed collection is not a subset of this bag')
self.discard_all(other)
[docs] def clear(self):
"""Remove all elements from this bag."""
self._dict = dict()
self._size = 0
# In-place operations
__ior__ = Bag._ior
__iand__ = Bag._iand
__ixor__ = Bag._ixor
__isub__ = Bag._isub
__iadd__ = Bag._iadd
[docs]class frozenbag(Bag, Hashable):
"""frozenbag is an immutable, hashable bag.
.. automethod:: __init__
"""
[docs] def __hash__(self):
"""Compute the hash value of a frozenbag."""
if not hasattr(self, '_hash_value'):
self._hash_value = Set._hash(self)
return self._hash_value
