""" Utility functions and classes """


def levenshtein_distance(s1, s2):
    """Compute the Levenshtein distance (edit distance) between two strings

    Shamelessly stolen from https://stackoverflow.com/a/32558749"""
    if len(s1) > len(s2):
        s1, s2 = s2, s1

    distances = range(len(s1) + 1)
    for i2, c2 in enumerate(s2):
        distances_ = [i2 + 1]
        for i1, c1 in enumerate(s1):
            if c1 == c2:
                distances_.append(distances[i1])
            else:
                distances_.append(
                    1 + min((distances[i1], distances[i1 + 1], distances_[-1]))
                )
        distances = distances_
    return distances[-1]


class UnionFind:
    """A union-find implementation"""

    parent_of: list[int]
    _group_size: list[int]

    def __init__(self, elt_count: int):
        self.parent_of = list(range(elt_count))
        self._group_size = [1] * elt_count

    def root(self, elt: int) -> int:
        """Find the element representing :elt: (root of component)

        Compresses paths along the way"""
        if self.parent_of[elt] == elt:
            return elt
        self.parent_of[elt] = self.root(self.parent_of[elt])
        return self.parent_of[elt]

    def union(self, elt1: int, elt2: int) -> None:
        """Unites two components"""
        elt1 = self.root(elt1)
        elt2 = self.root(elt2)
        if elt1 == elt2:
            return
        if self._group_size[elt1] > self._group_size[elt2]:
            self.union(elt2, elt1)
        else:
            self._group_size[elt2] += self._group_size[elt1]
            self._group_size[elt1] = 0
            self.parent_of[self.root(elt1)] = self.root(elt2)

    def group_size(self, elt: int) -> int:
        """Get the number of elements in the component of :elt:"""
        return self._group_size[self.root(elt)]