grouper.py

import typing

from typing import Generic, List, Optional, Set


T = typing.TypeVar('T')


class Grouper(Generic[T]):
    def __init__(self) -> None:
        self.groups: List[Set[T]] = []

    def _new_group(self) -> Set[T]:
        group: Set[T] = set()
        self.groups.append(group)
        return group

    def add(self, *values: T) -> None:
        # build group list
        groups = list(filter(None, [
            self.group_of(value) for value in values
        ]))
        if not groups:
            groups.append(self._new_group())
        # merge groups into the first
        target = groups[0]
        for merge_group in groups[1:]:
            if merge_group is not target:
                self.groups.remove(merge_group)
                target.update(merge_group)
        # add the values to the group
        target.update(values)

    def group_of(self, value: T) -> Optional[Set[T]]:
        for group in self.groups:
            if value in group:
                return group
        return None