fast_set_similarity.py

from bisect import bisect_left, bisect_right
from collections import defaultdict
from math import floor, sqrt


def iter_similar_sets(a, b, metric="jaccard", min_similarity=0.0):
    """
    Yield pairs of dictionary keys whose set similarity meets a threshold.

    Parameters
    ----------
    a, b : mapping
        Dictionaries whose keys identify sets and whose values are iterables of
        set members. Values are converted to sets internally, so duplicate
        members are ignored.
    metric : {"jaccard", "cosine", "dice"}, default="jaccard"
        Similarity metric to use.

        - `"jaccard"`: |A & B| / |A | B|
        - `"cosine"`: |A & B| / sqrt(|A| * |B|)
        - `"dice"`: 2 * |A & B| / (|A| + |B|)

    min_similarity : float, default=0.0
        Minimum similarity required for a pair to be yielded. This optimized
        implementation requires a positive threshold, because a threshold of
        zero would require yielding every pair, including disjoint pairs.

    Yields
    ------
    tuple
        Tuples of `(key_a, key_b, similarity)` for every pair whose similarity
        is greater than or equal to `min_similarity`.

    Raises
    ------
    ValueError
        If `metric` is not `"jaccard"`, `"cosine"`, or `"dice"`, or if
        `min_similarity` is outside `(0, 1]`.

    Notes
    -----
    The function builds an inverted index over `b` and only evaluates
    candidate pairs that share at least one member. Posting lists are sorted by
    set size, allowing candidate pairs that cannot possibly reach the threshold
    to be skipped using size bounds. When `a is b`, each unordered pair is
    considered once and self-pairs are skipped.
    """
    if not 0 < min_similarity <= 1:
        raise ValueError("min_similarity must be greater than 0 and at most 1")

    same_dict = a is b
    a_items = [(key, set(values)) for key, values in a.items()]
    b_items = (
        a_items if same_dict else [(key, set(values)) for key, values in b.items()]
    )

    b_keys = [key for key, _ in b_items]
    b_sizes = [len(values) for _, values in b_items]
    b_sqrt_sizes = [sqrt(size) for size in b_sizes]

    raw_index = defaultdict(list)
    for b_id, (_, values) in enumerate(b_items):
        for value in values:
            raw_index[value].append(b_id)

    index = {}
    for value, b_ids in raw_index.items():
        b_ids.sort(key=b_sizes.__getitem__)
        index[value] = ([b_sizes[b_id] for b_id in b_ids], b_ids)

    threshold = float(min_similarity)
    threshold_sq = threshold * threshold

    for a_id, (key_a, values_a) in enumerate(a_items):
        size_a = len(values_a)
        if size_a == 0:
            continue

        match metric:
            case "jaccard":
                min_size_b = threshold * size_a
                max_size_b = size_a / threshold
            case "cosine":
                min_size_b = threshold_sq * size_a
                max_size_b = size_a / threshold_sq
            case "dice":
                min_size_b = threshold * size_a / (2 - threshold)
                max_size_b = (2 - threshold) * size_a / threshold
            case _:
                raise ValueError("metric must be 'jaccard', 'cosine', or 'dice'")

        counts = {}

        for value in values_a:
            posting = index.get(value)
            if posting is None:
                continue

            posting_sizes, posting_ids = posting
            left = bisect_left(posting_sizes, min_size_b)
            right = bisect_right(posting_sizes, floor(max_size_b))

            for b_id in posting_ids[left:right]:
                if same_dict and b_id <= a_id:
                    continue
                counts[b_id] = counts.get(b_id, 0) + 1

        match metric:
            case "jaccard":
                for b_id, intersection in counts.items():
                    size_b = b_sizes[b_id]
                    if intersection * (1 + threshold) >= threshold * (size_a + size_b):
                        yield (
                            key_a,
                            b_keys[b_id],
                            intersection / (size_a + size_b - intersection),
                        )

            case "cosine":
                sqrt_size_a = sqrt(size_a)

                for b_id, intersection in counts.items():
                    denominator = sqrt_size_a * b_sqrt_sizes[b_id]
                    if intersection >= threshold * denominator:
                        yield key_a, b_keys[b_id], intersection / denominator

            case "dice":
                for b_id, intersection in counts.items():
                    size_b = b_sizes[b_id]
                    if 2 * intersection >= threshold * (size_a + size_b):
                        yield key_a, b_keys[b_id], 2 * intersection / (size_a + size_b)