karaokatalog/karaokatalog/util/get_equivalence_classes.py
2025-05-21 10:17:50 +02:00

35 lines
1.5 KiB
Python

from collections.abc import Iterable, Sequence, Callable
type EquivalenceClass[T] = Sequence[T]
def get_equivalence_classes[T](
items: Iterable[T], is_equivalent: Callable[[T, T], bool]
) -> Sequence[EquivalenceClass[T]]:
"""
Partition an iterable of items into equivalence classes under the given equivalence relation.
The order of `items` is kept within the equivalence classes.
>>> from math import floor
>>> get_equivalence_classes([1.0, 2.1, 1.1, 1.0, 2.2, 2.1, 3.3], lambda a, b: floor(a) == floor(b))
((1.0, 1.1, 1.0), (2.1, 2.2, 2.1), (3.3,))
>>> get_equivalence_classes(range(15), lambda a, b: a % 3 == b % 3)
((0, 3, 6, 9, 12), (1, 4, 7, 10, 13), (2, 5, 8, 11, 14))
"""
equivalence_classes: list[list[T]] = []
for item in items:
# Check if the item belongs into an already existing equivalence class
for equivalence_class in equivalence_classes:
# Pick an arbitrary representative of the equivalence class (we are allowed to do this
# because it is an equivalence class)
representative = equivalence_class[0]
if is_equivalent(item, representative):
equivalence_class.append(item)
break
else:
# This item forms a new equivalence class, so we create one containing only the item and append it
equivalence_classes.append([item])
return tuple(tuple(equivalence_class) for equivalence_class in equivalence_classes)