from typing import List
from typing import Set


def sparse_closure(buckets_by_document: List[List[int]], documents_by_bucket: List[List[int]]):
    """
    Computes transitive closure via some BFS/DFS mixture on a sparse binary matrix of minhash buckets. Imagine a
    (sparse) matrix where rows represent documents, columns represent buckets and binary values indicate whether a
    document is part of a bucket. This method accepts such a matrix in a very condensed form, namely a list of list of
    bucket indices (document -> buckets mapping) and a list of list of document indices (bucket -> documents mapping).
    Returns the closure / clustering as a list of list of document indices.
    """
    docs_not_visited = set(range(len(buckets_by_document)))
    clustering = []

    def get_buckets(rows: Set) -> Set:
        return set(bucket for row in rows for bucket in buckets_by_document[row])
    def get_documents(col) -> Set:
        return set(documents_by_bucket[col])

    while docs_not_visited:
        seed_doc = docs_not_visited.pop()
        docs_visited = {seed_doc}                   # this is the current cluster in the making
        buckets_visited = set()
        bucket_queue = get_buckets({seed_doc})      # find buckets this document is part of (horizontal matrix lookup)
        while bucket_queue:
            bucket = bucket_queue.pop()
            buckets_visited.add(bucket)
            docs = get_documents(bucket)            # find documents in this bucket (vertical matrix lookup)
            connected_docs_not_yet_visited = docs - docs_visited
            if connected_docs_not_yet_visited:
                bucket_queue |= get_buckets(connected_docs_not_yet_visited) - buckets_visited
                docs_visited |= docs
        docs_not_visited -= docs_visited
        clustering.append(list(docs_visited))

    return clustering