# -*- coding: utf-8 -*-
"""An implementation of a multiset."""
from typing import (Generic, TypeVar, Hashable, Mapping as MappingType, Union, Optional, Iterable as IterableType,
Type, ItemsView, KeysView, ValuesView, MutableMapping as MutableMappingType, AbstractSet as SetType)
from collections import defaultdict
from collections.abc import Iterable, Mapping, MutableMapping, Set, Sized, Container
from itertools import chain, repeat, starmap
_sequence_types = (tuple, list, range, set, frozenset, str)
_iter_types = (type(iter([])), type((lambda: (yield))()))
_all_basic_types = _sequence_types + _iter_types + (dict, )
__all__ = ['BaseMultiset', 'Multiset', 'FrozenMultiset']
_TElement = TypeVar('_TElement', bound=Hashable)
_OtherType = Union[IterableType[_TElement], MappingType[_TElement, int]]
_Self = TypeVar('_Self', bound='BaseMultiset')
[docs]class BaseMultiset(MappingType[_TElement, int], Generic[_TElement]):
"""A multiset implementation.
A multiset is similar to the builtin :class:`set`, but elements can occur multiple times in the multiset.
It is also similar to a :class:`list` without ordering of the values and hence no index-based operations.
The multiset internally uses a :class:`dict` for storage where the key is the element and the value its
multiplicity. It supports all operations that the :class:`set` supports.
In contrast to the builtin :class:`collections.Counter`, no negative counts are allowed, elements with
zero counts are removed from the :class:`dict`, and set operations are supported.
The multiset comes in two variants, `Multiset` and `FrozenMultiset` which correspond to the `set` and
`frozenset` classes, respectively.
.. warning::
You cannot instantiate this class directly. Use one of its variants instead.
:see: https://en.wikipedia.org/wiki/Multiset
"""
__slots__ = ('_elements', '_total')
[docs] def __init__(self, iterable: Optional[_OtherType] = None):
r"""Create a new, empty Multiset object.
And if given, initialize with elements from input iterable.
Or, initialize from a mapping of elements to their multiplicity.
Example:
>>> ms = Multiset() # a new, empty multiset
>>> ms = Multiset('abc') # a new multiset from an iterable
>>> ms = Multiset({'a': 4, 'b': 2}) # a new multiset from a mapping
Args:
iterable:
An optional iterable of elements or mapping of elements to multiplicity to
initialize the multiset from.
"""
if isinstance(iterable, BaseMultiset):
self._elements = iterable._elements.copy()
self._total = iterable._total
else:
self._elements = _elements = defaultdict(int)
_total = 0
if iterable is not None:
if isinstance(iterable, _sequence_types):
for element in iterable:
_elements[element] += 1
_total = len(iterable)
elif isinstance(iterable, dict):
for element, multiplicity in iterable.items():
if multiplicity > 0:
_elements[element] = multiplicity
_total += multiplicity
elif isinstance(iterable, _iter_types):
for element in iterable:
_elements[element] += 1
_total += 1
elif isinstance(iterable, Mapping):
for element, multiplicity in iterable.items():
if multiplicity > 0:
_elements[element] = multiplicity
_total += multiplicity
elif isinstance(iterable, Sized):
for element in iterable:
_elements[element] += 1
_total = len(iterable)
else:
for element in iterable:
_elements[element] += 1
_total += 1
self._total = _total
def __new__(cls, iterable=None):
if cls is BaseMultiset:
raise TypeError("Cannot instantiate BaseMultiset directly, use either Multiset or FrozenMultiset.")
return super(BaseMultiset, cls).__new__(cls)
def __contains__(self, element: object) -> bool:
return element in self._elements
def __getitem__(self, element: _TElement) -> int:
"""The multiplicity of an element or zero if it is not in the multiset."""
return self._elements.get(element, 0)
def __str__(self) -> str:
return '{%s}' % ', '.join(map(str, self.__iter__()))
def __repr__(self) -> str:
items = ', '.join('%r: %r' % item for item in self._elements.items())
return '%s({%s})' % (self.__class__.__name__, items)
def __len__(self) -> int:
"""Returns the total number of elements in the multiset.
Note that this is equivalent to the sum of the multiplicities:
>>> ms = Multiset('aab')
>>> len(ms)
3
>>> sum(ms.multiplicities())
3
If you need the total number of distinct elements, use either the :meth:`distinct_elements` method:
>>> len(ms.distinct_elements())
2
or convert to a :class:`set`:
>>> len(set(ms))
2
"""
return self._total
def __bool__(self) -> bool:
return self._total > 0
def __iter__(self) -> IterableType[_TElement]:
return chain.from_iterable(starmap(repeat, self._elements.items()))
[docs] def isdisjoint(self, other: _OtherType) -> bool:
r"""Return True if the set has no elements in common with other.
Sets are disjoint iff their intersection is the empty set.
>>> ms = Multiset('aab')
>>> ms.isdisjoint('bc')
False
>>> ms.isdisjoint(Multiset('ccd'))
True
Args:
other: The other set to check disjointedness. Can also be an :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
"""
if isinstance(other, _sequence_types + (BaseMultiset, )):
pass
elif not isinstance(other, Container):
other = self._as_multiset(other)
return all(element not in other for element in self._elements.keys())
[docs] def difference(self: _Self, *others: _OtherType) -> _Self:
r"""Return a new multiset with all elements from the others removed.
>>> ms = Multiset('aab')
>>> sorted(ms.difference('bc'))
['a', 'a']
You can also use the ``-`` operator for the same effect. However, the operator version
will only accept a set as other operator, not any iterable, to avoid errors.
>>> ms = Multiset('aabbbc')
>>> sorted(ms - Multiset('abd'))
['a', 'b', 'b', 'c']
For a variant of the operation which modifies the multiset in place see
:meth:`difference_update`.
Args:
others: The other sets to remove from the multiset. Can also be any :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
Returns:
The resulting difference multiset.
"""
result = self.__copy__()
_elements = result._elements
_total = result._total
for other in map(self._as_multiset, others):
for element, multiplicity in other.items():
if element in _elements:
old_multiplicity = _elements[element]
new_multiplicity = old_multiplicity - multiplicity
if new_multiplicity > 0:
_elements[element] = new_multiplicity
_total -= multiplicity
else:
del _elements[element]
_total -= old_multiplicity
result._total = _total
return result
def __sub__(self: _Self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> _Self:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
return self.difference(other)
def __rsub__(self: _Self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> _Self:
if not isinstance(other, (Set, BaseMultiset)):
return NotImplemented
return self._as_multiset(other).difference(self)
[docs] def union(self: _Self, *others: _OtherType) -> _Self:
r"""Return a new multiset with all elements from the multiset and the others with maximal multiplicities.
>>> ms = Multiset('aab')
>>> sorted(ms.union('bc'))
['a', 'a', 'b', 'c']
You can also use the ``|`` operator for the same effect. However, the operator version
will only accept a set as other operator, not any iterable, to avoid errors.
>>> ms = Multiset('aab')
>>> sorted(ms | Multiset('aaa'))
['a', 'a', 'a', 'b']
For a variant of the operation which modifies the multiset in place see
:meth:`union_update`.
Args:
*others: The other sets to union the multiset with. Can also be any :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
Returns:
The multiset resulting from the union.
"""
result = self.__copy__()
_elements = result._elements
_total = result._total
for other in map(self._as_mapping, others):
for element, multiplicity in other.items():
old_multiplicity = _elements.get(element, 0)
if multiplicity > old_multiplicity:
_elements[element] = multiplicity
_total += multiplicity - old_multiplicity
result._total = _total
return result
def __or__(self: _Self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> _Self:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
return self.union(other)
__ror__ = __or__
[docs] def combine(self: _Self, *others: _OtherType) -> _Self:
r"""Return a new multiset with all elements from the multiset and the others with their multiplicities summed up.
>>> ms = Multiset('aab')
>>> sorted(ms.combine('bc'))
['a', 'a', 'b', 'b', 'c']
You can also use the ``+`` operator for the same effect. However, the operator version
will only accept a set as other operator, not any iterable, to avoid errors.
>>> ms = Multiset('aab')
>>> sorted(ms + Multiset('a'))
['a', 'a', 'a', 'b']
For a variant of the operation which modifies the multiset in place see
:meth:`update`.
Args:
others: The other sets to add to the multiset. Can also be any :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
Returns:
The multiset resulting from the addition of the sets.
"""
result = self.__copy__()
_elements = result._elements
_total = result._total
for other in map(self._as_mapping, others):
for element, multiplicity in other.items():
old_multiplicity = _elements.get(element, 0)
new_multiplicity = old_multiplicity + multiplicity
if old_multiplicity > 0 and new_multiplicity <= 0:
del _elements[element]
_total -= old_multiplicity
elif new_multiplicity > 0:
_elements[element] = new_multiplicity
_total += multiplicity
result._total = _total
return result
def __add__(self: _Self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> _Self:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
return self.combine(other)
__radd__ = __add__
[docs] def intersection(self: _Self, *others: _OtherType) -> _Self:
r"""Return a new multiset with elements common to the multiset and all others.
>>> ms = Multiset('aab')
>>> sorted(ms.intersection('abc'))
['a', 'b']
You can also use the ``&`` operator for the same effect. However, the operator version
will only accept a set as other operator, not any iterable, to avoid errors.
>>> ms = Multiset('aab')
>>> sorted(ms & Multiset('aaac'))
['a', 'a']
For a variant of the operation which modifies the multiset in place see
:meth:`intersection_update`.
Args:
others: The other sets intersect with the multiset. Can also be any :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
Returns:
The multiset resulting from the intersection of the sets.
"""
result = self.__copy__()
_elements = result._elements
_total = result._total
for other in map(self._as_mapping, others):
for element, multiplicity in list(_elements.items()):
new_multiplicity = other.get(element, 0)
if new_multiplicity < multiplicity:
if new_multiplicity > 0:
_elements[element] = new_multiplicity
_total -= multiplicity - new_multiplicity
else:
del _elements[element]
_total -= multiplicity
result._total = _total
return result
def __and__(self: _Self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> _Self:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
return self.intersection(other)
__rand__ = __and__
[docs] def symmetric_difference(self: _Self, other: _OtherType) -> _Self:
r"""Return a new set with elements in either the set or other but not both.
>>> ms = Multiset('aab')
>>> sorted(ms.symmetric_difference('abc'))
['a', 'c']
You can also use the ``^`` operator for the same effect. However, the operator version
will only accept a set as other operator, not any iterable, to avoid errors.
>>> ms = Multiset('aab')
>>> sorted(ms ^ Multiset('aaac'))
['a', 'b', 'c']
For a variant of the operation which modifies the multiset in place see
:meth:`symmetric_difference_update`.
Args:
other: The other set to take the symmetric difference with. Can also be any :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
Returns:
The resulting symmetric difference multiset.
"""
other = self._as_multiset(other)
result = self.__class__()
_total = 0
_elements = result._elements
self_elements = self._elements
other_elements = other._elements
dist_elements = set(self_elements.keys()) | set(other_elements.keys())
for element in dist_elements:
multiplicity = self_elements.get(element, 0)
other_multiplicity = other_elements.get(element, 0)
new_multiplicity = (multiplicity - other_multiplicity
if multiplicity > other_multiplicity else other_multiplicity - multiplicity)
_total += new_multiplicity
if new_multiplicity > 0:
_elements[element] = new_multiplicity
result._total = _total
return result
def __xor__(self: _Self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> _Self:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
return self.symmetric_difference(other)
__rxor__ = __xor__
[docs] def times(self: _Self, factor: int) -> _Self:
"""Return a new set with each element's multiplicity multiplied with the given scalar factor.
>>> ms = Multiset('aab')
>>> sorted(ms.times(2))
['a', 'a', 'a', 'a', 'b', 'b']
You can also use the ``*`` operator for the same effect:
>>> sorted(ms * 3)
['a', 'a', 'a', 'a', 'a', 'a', 'b', 'b', 'b']
For a variant of the operation which modifies the multiset in place see
:meth:`times_update`.
Args:
factor: The factor to multiply each multiplicity with.
"""
if factor == 0:
return self.__class__()
if factor < 0:
raise ValueError('The factor must no be negative.')
result = self.__copy__()
_elements = result._elements
for element in _elements:
_elements[element] *= factor
result._total *= factor
return result
def __mul__(self: _Self, factor: int) -> _Self:
if not isinstance(factor, int):
return NotImplemented
return self.times(factor)
__rmul__ = __mul__
def _issubset(self, other: _OtherType, strict: bool) -> bool:
other = self._as_multiset(other)
self_len = self._total
other_len = len(other)
if self_len > other_len:
return False
if self_len == other_len and strict:
return False
return all(multiplicity <= other[element] for element, multiplicity in self.items())
[docs] def issubset(self, other: _OtherType) -> bool:
"""Return True iff this set is a subset of the other.
>>> Multiset('ab').issubset('aabc')
True
>>> Multiset('aabb').issubset(Multiset('aabc'))
False
You can also use the ``<=`` operator for this comparison:
>>> Multiset('ab') <= Multiset('ab')
True
When using the ``<`` operator for comparison, the sets are checked
to be unequal in addition:
>>> Multiset('ab') < Multiset('ab')
False
Args:
other: The potential superset of the multiset to be checked.
Returns:
True iff this set is a subset of the other.
"""
return self._issubset(other, False)
def __le__(self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> bool:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
return self._issubset(other, False)
def __lt__(self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> bool:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
return self._issubset(other, True)
def _issuperset(self, other: _OtherType, strict: bool) -> bool:
other = self._as_multiset(other)
other_len = len(other)
if len(self) < other_len:
return False
if len(self) == other_len and strict:
return False
for element, multiplicity in other.items():
if self[element] < multiplicity:
return False
return True
[docs] def issuperset(self, other: _OtherType) -> bool:
"""Return True iff this multiset is a superset of the other.
>>> Multiset('aabc').issuperset('ab')
True
>>> Multiset('aabc').issuperset(Multiset('abcc'))
False
You can also use the ``>=`` operator for this comparison:
>>> Multiset('ab') >= Multiset('ab')
True
When using the ``>`` operator for comparison, the sets are checked
to be unequal in addition:
>>> Multiset('ab') > Multiset('ab')
False
Args:
other: The potential subset of the multiset to be checked.
Returns:
True iff this set is a subset of the other.
"""
return self._issuperset(other, False)
def __ge__(self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> bool:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
return self._issuperset(other, False)
def __gt__(self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> bool:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
return self._issuperset(other, True)
def __eq__(self, other: object) -> bool:
if isinstance(other, BaseMultiset):
return self._total == other._total and self._elements == other._elements
if isinstance(other, (set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
if self._total != len(other):
return False
return self._issubset(other, False)
def __ne__(self, other: object) -> bool:
if isinstance(other, BaseMultiset):
return self._total != other._total or self._elements != other._elements
if isinstance(other, (set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
if self._total != len(other):
return True
return not self._issubset(other, False)
[docs] def get(self, element: _TElement, default: int) -> int:
"""Return the multiplicity for *element* if it is in the multiset, else *default*.
Makes the *default* argument of the original :meth:`dict.get` non-optional.
Args:
element: The element of which to get the multiplicity.
default: The default value to return if the element if not in the multiset.
Returns:
The multiplicity for *element* if it is in the multiset, else *default*.
"""
return self._elements.get(element, default)
[docs] @classmethod
def from_elements(cls: Type[_Self], elements: IterableType[_TElement], multiplicity: int) -> _Self:
"""Create a new multiset with the given *elements* and each multiplicity set to *multiplicity*.
Uses :meth:`dict.fromkeys` internally.
Args:
elements: The element for the new multiset.
multiplicity: The multiplicity for all elements.
Returns:
The new multiset.
"""
return cls(dict.fromkeys(elements, multiplicity))
[docs] def copy(self: _Self) -> _Self:
"""Return a shallow copy of the multiset."""
return self.__class__(self)
__copy__ = copy
[docs] def items(self) -> ItemsView[_TElement, int]:
return self._elements.items()
[docs] def distinct_elements(self) -> KeysView[_TElement]:
return self._elements.keys()
[docs] def multiplicities(self) -> ValuesView[int]:
return self._elements.values()
values = multiplicities
@classmethod
def _as_multiset(cls, other):
if isinstance(other, BaseMultiset):
return other
if isinstance(other, _all_basic_types):
pass
elif not isinstance(other, Iterable):
raise TypeError("'%s' object is not iterable" % type(other)) # pragma: no cover
return cls(other)
@staticmethod
def _as_mapping(iterable):
if isinstance(iterable, BaseMultiset):
return iterable._elements
if isinstance(iterable, dict):
return iterable
if isinstance(iterable, _all_basic_types):
pass # create dictionary below
elif isinstance(iterable, Mapping):
return iterable
elif not isinstance(iterable, Iterable):
raise TypeError("'%s' object is not iterable" % type(iterable))
mapping = dict()
for element in iterable:
if element in mapping:
mapping[element] += 1
else:
mapping[element] = 1
return mapping
def __getstate__(self):
return self._total, self._elements
def __setstate__(self, state):
self._total, self._elements = state
[docs]class Multiset(BaseMultiset[_TElement], MutableMappingType[_TElement, int], Generic[_TElement]):
"""The mutable multiset variant."""
__slots__ = ()
def __setitem__(self, element: _TElement, multiplicity: int):
"""Set the element's multiplicity.
This will remove the element if the multiplicity is less than or equal to zero.
'"""
if not isinstance(multiplicity, int):
raise TypeError('multiplicity must be an integer')
_elements = self._elements
if element in _elements:
old_multiplicity = _elements[element]
if multiplicity > 0:
_elements[element] = multiplicity
self._total += multiplicity - old_multiplicity
else:
del _elements[element]
self._total -= old_multiplicity
elif multiplicity > 0:
_elements[element] = multiplicity
self._total += multiplicity
def __delitem__(self, element: _TElement):
_elements = self._elements
if element in _elements:
self._total -= _elements[element]
del _elements[element]
else:
raise KeyError("Could not delete {!r} from the multiset, because it is not in it.".format(element))
[docs] def update(self, *others: _OtherType, **kwargs):
r"""Like :meth:`dict.update` but add multiplicities instead of replacing them.
>>> ms = Multiset('aab')
>>> ms.update('abc')
>>> sorted(ms)
['a', 'a', 'a', 'b', 'b', 'c']
Note that the operator ``+=`` is equivalent to :meth:`update`, except that the operator will only
accept sets to avoid accidental errors.
>>> ms += Multiset('bc')
>>> sorted(ms)
['a', 'a', 'a', 'b', 'b', 'b', 'c', 'c']
For a variant of the operation which does not modify the multiset, but returns a new
multiset instead see :meth:`combine`.
Any keyword arguments are also added to the multiset:
>>> ms = Multiset('ab')
>>> ms.update(a=1, e=2)
>>> sorted(ms)
['a', 'a', 'b', 'e', 'e']
Args:
others: The other sets to add to this multiset. Can also be any :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
kwargs: Additional pairs of values and multiplicities to be added to multiset.
"""
_elements = self._elements
_total = self._total
for other in map(self._as_mapping, others + (kwargs, )):
for element, multiplicity in other.items():
if multiplicity > 0:
old_multiplicity = _elements.get(element, 0)
_elements[element] = multiplicity + old_multiplicity
_total += multiplicity
self._total = _total
[docs] def union_update(self, *others: _OtherType):
r"""Update the multiset, adding elements from all others using the maximum multiplicity.
>>> ms = Multiset('aab')
>>> ms.union_update('bc')
>>> sorted(ms)
['a', 'a', 'b', 'c']
You can also use the ``|=`` operator for the same effect. However, the operator version
will only accept a set as other operator, not any iterable, to avoid errors.
>>> ms = Multiset('aab')
>>> ms |= Multiset('bccd')
>>> sorted(ms)
['a', 'a', 'b', 'c', 'c', 'd']
For a variant of the operation which does not modify the multiset, but returns a new
multiset instead see :meth:`union`.
Args:
others: The other sets to union this multiset with. Can also be any :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
"""
_elements = self._elements
_total = self._total
for other in map(self._as_mapping, others):
for element, multiplicity in other.items():
old_multiplicity = _elements.get(element, 0)
if multiplicity > old_multiplicity:
_elements[element] = multiplicity
_total += multiplicity - old_multiplicity
self._total = _total
def __ior__(self: _Self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> _Self:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
self.union_update(other)
return self
[docs] def intersection_update(self, *others: _OtherType):
r"""Update the multiset, keeping only elements found in it and all others.
>>> ms = Multiset('aab')
>>> ms.intersection_update('bc')
>>> sorted(ms)
['b']
You can also use the ``&=`` operator for the same effect. However, the operator version
will only accept a set as other operator, not any iterable, to avoid errors.
>>> ms = Multiset('aabc')
>>> ms &= Multiset('abbd')
>>> sorted(ms)
['a', 'b']
For a variant of the operation which does not modify the multiset, but returns a new
multiset instead see :meth:`intersection`.
Args:
others: The other sets to intersect this multiset with. Can also be any :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
"""
for other in map(self._as_mapping, others):
for element, current_count in list(self.items()):
multiplicity = other.get(element, 0)
if multiplicity < current_count:
self[element] = multiplicity
def __iand__(self: _Self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> _Self:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
self.intersection_update(other)
return self
[docs] def difference_update(self, *others: _OtherType):
r"""Remove all elements contained the others from this multiset.
>>> ms = Multiset('aab')
>>> ms.difference_update('abc')
>>> sorted(ms)
['a']
You can also use the ``-=`` operator for the same effect. However, the operator version
will only accept a set as other operator, not any iterable, to avoid errors.
>>> ms = Multiset('aabbbc')
>>> ms -= Multiset('abd')
>>> sorted(ms)
['a', 'b', 'b', 'c']
For a variant of the operation which does not modify the multiset, but returns a new
multiset instead see :meth:`difference`.
Args:
others: The other sets to remove from this multiset. Can also be any :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
"""
for other in map(self._as_multiset, others):
for element, multiplicity in other.items():
self.discard(element, multiplicity)
def __isub__(self: _Self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> _Self:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
self.difference_update(other)
return self
[docs] def symmetric_difference_update(self, other: _OtherType):
r"""Update the multiset to contain only elements in either this multiset or the other but not both.
>>> ms = Multiset('aab')
>>> ms.symmetric_difference_update('abc')
>>> sorted(ms)
['a', 'c']
You can also use the ``^=`` operator for the same effect. However, the operator version
will only accept a set as other operator, not any iterable, to avoid errors.
>>> ms = Multiset('aabbbc')
>>> ms ^= Multiset('abd')
>>> sorted(ms)
['a', 'b', 'b', 'c', 'd']
For a variant of the operation which does not modify the multiset, but returns a new
multiset instead see :meth:`symmetric_difference`.
Args:
other: The other set to take the symmetric difference with. Can also be any :class:`~typing.Iterable`\[~T]
or :class:`~typing.Mapping`\[~T, :class:`int`] which are then converted to :class:`Multiset`\[~T].
"""
other = self._as_multiset(other)
elements = set(self.distinct_elements()) | set(other.distinct_elements())
for element in elements:
multiplicity = self[element]
other_count = other[element]
self[element] = (multiplicity - other_count if multiplicity > other_count else other_count - multiplicity)
def __ixor__(self: _Self, other: Union[SetType[_TElement], 'BaseMultiset[_TElement]']) -> _Self:
if isinstance(other, (BaseMultiset, set, frozenset)):
pass
elif not isinstance(other, Set):
return NotImplemented
self.symmetric_difference_update(other)
return self
[docs] def times_update(self, factor: int):
"""Update each this multiset by multiplying each element's multiplicity with the given scalar factor.
>>> ms = Multiset('aab')
>>> ms.times_update(2)
>>> sorted(ms)
['a', 'a', 'a', 'a', 'b', 'b']
You can also use the ``*=`` operator for the same effect:
>>> ms = Multiset('ac')
>>> ms *= 3
>>> sorted(ms)
['a', 'a', 'a', 'c', 'c', 'c']
For a variant of the operation which does not modify the multiset, but returns a new
multiset instead see :meth:`times`.
Args:
factor: The factor to multiply each multiplicity with.
"""
if factor < 0:
raise ValueError("The factor must not be negative.")
elif factor == 0:
self.clear()
else:
_elements = self._elements
for element in _elements:
_elements[element] *= factor
self._total *= factor
def __imul__(self: _Self, factor: int) -> _Self:
if not isinstance(factor, int):
raise TypeError("factor must be an integer.")
self.times_update(factor)
return self
[docs] def add(self, element: _TElement, multiplicity: int = 1):
"""Adds an element to the multiset.
>>> ms = Multiset()
>>> ms.add('a')
>>> sorted(ms)
['a']
An optional multiplicity can be specified to define how many of the element are added:
>>> ms.add('b', 2)
>>> sorted(ms)
['a', 'b', 'b']
This extends the :meth:`MutableSet.add` signature to allow specifying the multiplicity.
Args:
element:
The element to add to the multiset.
multiplicity:
The multiplicity i.e. count of elements to add.
"""
if multiplicity < 1:
raise ValueError("Multiplicity must be positive")
self._elements[element] += multiplicity
self._total += multiplicity
[docs] def remove(self, element: _TElement, multiplicity: Optional[int] = None) -> int:
"""Removes an element from the multiset.
If no multiplicity is specified, the element is completely removed from the multiset:
>>> ms = Multiset('aabbbc')
>>> ms.remove('a')
2
>>> sorted(ms)
['b', 'b', 'b', 'c']
If the multiplicity is given, it is subtracted from the element's multiplicity in the multiset:
>>> ms.remove('b', 2)
3
>>> sorted(ms)
['b', 'c']
It is not an error to remove more elements than are in the set:
>>> ms.remove('b', 2)
1
>>> sorted(ms)
['c']
This extends the :meth:`MutableSet.remove` signature to allow specifying the multiplicity.
Args:
element:
The element to remove from the multiset.
multiplicity:
An optional multiplicity i.e. count of elements to remove.
Returns:
The multiplicity of the element in the multiset before
the removal.
Raises:
KeyError: if the element is not contained in the set. Use :meth:`discard` if
you do not want an exception to be raised.
"""
_elements = self._elements
if element not in _elements:
raise KeyError
old_multiplicity = _elements.get(element, 0)
if multiplicity is None or multiplicity >= old_multiplicity:
del _elements[element]
self._total -= old_multiplicity
elif multiplicity < 0:
raise ValueError("Multiplicity must be not be negative")
elif multiplicity > 0:
_elements[element] -= multiplicity
self._total -= multiplicity
return old_multiplicity
[docs] def discard(self, element: _TElement, multiplicity: Optional[int] = None) -> int:
"""Removes the `element` from the multiset.
If multiplicity is ``None``, all occurrences of the element are removed:
>>> ms = Multiset('aab')
>>> ms.discard('a')
2
>>> sorted(ms)
['b']
Otherwise, the multiplicity is subtracted from the one in the multiset and the
old multiplicity is removed:
>>> ms = Multiset('aab')
>>> ms.discard('a', 1)
2
>>> sorted(ms)
['a', 'b']
In contrast to :meth:`remove`, this does not raise an error if the
element is not in the multiset:
>>> ms = Multiset('a')
>>> ms.discard('b')
0
>>> sorted(ms)
['a']
It is also not an error to remove more elements than are in the set:
>>> ms.remove('a', 2)
1
>>> sorted(ms)
[]
Args:
element:
The element to remove from the multiset.
multiplicity:
An optional multiplicity i.e. count of elements to remove.
Returns:
The multiplicity of the element in the multiset before
the removal.
"""
_elements = self._elements
if element in _elements:
old_multiplicity = _elements[element]
if multiplicity is None or multiplicity >= old_multiplicity:
del _elements[element]
self._total -= old_multiplicity
elif multiplicity < 0:
raise ValueError("Multiplicity must not be negative")
elif multiplicity > 0:
_elements[element] -= multiplicity
self._total -= multiplicity
return old_multiplicity
else:
return 0
[docs] def pop(self, element: _TElement, default: int) -> int:
"""If *element* is in the multiset, remove it and return its multiplicity, else return *default*.
Makes the *default* argument of the original :meth:`dict.pop` non-optional.
Args:
element: The element which is removed.
default: The default value to return if the element if not in the multiset.
Returns:
The multiplicity for *element* if it is in the multiset, else *default*.
"""
rm_size = self._elements.get(element)
if rm_size is not None:
self._total -= rm_size
return self._elements.pop(element, default)
[docs] def setdefault(self, element: _TElement, default: int) -> int:
"""If *element* is in the multiset, return its multiplicity.
Else add it with a multiplicity of *default* and return *default*.
Makes the *default* argument of the original :meth:`dict.setdefault` non-optional.
Args:
element: The element which is added if not already present.
default: The default multiplicity to add the element with if not in the multiset.
Returns:
The multiplicity for *element* if it is in the multiset, else *default*.
"""
mul = self._elements.get(element)
if mul is None:
if default < 1:
raise ValueError("Multiplicity must be positive")
self._total += default
return self._elements.setdefault(element, default)
[docs] def clear(self):
"""Empty the multiset."""
self._elements.clear()
self._total = 0
[docs]class FrozenMultiset(BaseMultiset[_TElement], Generic[_TElement]):
"""The frozen multiset variant that is immutable and hashable."""
__slots__ = ()
def __hash__(self):
return hash(frozenset(self._elements.items()))
Mapping.register(BaseMultiset) # type: ignore
MutableMapping.register(Multiset) # type: ignore
if __name__ == '__main__':
import doctest
doctest.testmod()