fast-unionfind

Disjoint sets (union-find) for Python at C speed, with no dependencies.

from fast_unionfind import unionfind

uf = unionfind(1_000_000)
uf.unite(3, 7)          # True: two sets became one
uf.unite(7, 3)          # False: already the same set
uf.connected(3, 7)      # True
uf[7] == uf.find(3)     # True

It is link-by-rank with path halving, so any sequence of m operations over n elements costs O(m α(n)) — effectively constant per operation. It has no runtime dependencies at all (not even numpy), it is compiled with Cython, and it is safe to query from many threads at once, including on free-threaded Python.

Why another union-find

Two reasons.

There was no C-speed union-find available as a stand-alone package. The fast implementations live inside larger libraries, as internals you cannot depend on.

Building one efficiently is less trivial than it looks. The union-finds inside two widely used clustering libraries, hdbscan and scikit-learn, had subtle performance bugs that made single-linkage labelling quadratic while every result stayed correct. We found and reported both: hdbscan#731 and scikit-learn#34626 (a fix is in review as #34872).

Install

pip install fast-unionfind

Building from source needs a C compiler and Cython.

Python API

unionfind(n) builds a union-find over the elements 0 … n-1 and picks the narrowest index width that fits: a UnionFind32 (4 bytes per element) up to 2³²−1 elements, a UnionFind64 (8 bytes) beyond that. Either class can also be built directly. UnionFind is their common base, for annotations and isinstance checks, and cannot be instantiated itself.

Elements are ids, not sequence positions. Anything outside range(n) raises IndexError, negative values included, so uf[-1] is an error rather than the last element.

Representatives change as sets merge, so compare them rather than store them.

There is no set count or grouping method, because both are one visible line on top of labels(), and writing them out keeps their O(n) cost in plain sight:

n_sets = len(set(uf.labels()))

groups = collections.defaultdict(list)
for element, root in enumerate(uf.labels()):
    groups[root].append(element)

labels() returns a buffer, so numpy.frombuffer(uf.labels(), dtype=...) wraps it without copying if you do have numpy.

Union-finds pickle exactly, ranks included, so a restored one links later unions the same way the original would have.

Cython API

The union-find itself is not a method. It is a pair of fused cdef inline functions defined in the .pxd, so a module that cimports them compiles its own copy, specialised for its index width, and the C compiler inlines them into your loop:

from libc.stdint cimport uint32_t
from fast_unionfind.core cimport UnionFind32, uf_find, uf_unite

def kruskal(Py_ssize_t n, uint32_t[::1] ei, uint32_t[::1] ej):
    cdef UnionFind32 uf = UnionFind32(n)      # owns the memory
    cdef uint32_t* parents = uf.parents       # plain struct reads:
    cdef unsigned char* ranks = uf.ranks      # the class is final
    cdef Py_ssize_t k, merged = 0
    with nogil:
        for k in range(ei.shape[0]):
            if uf_unite(parents, ranks, ei[k], ej[k]):
                merged += 1
    return merged

Both are fused over uint32_t and uint64_t. Use UnionFind64's parents for the wide one. They do no bounds checking, which is the point of this layer. Keep the owning UnionFind32/UnionFind64 alive for as long as you use its pointers.

The .pxd files ship in the wheel, so an installed package is enough for Cython to resolve the cimport. To be explicit anyway, add fast_unionfind.get_include() to your extension's include_dirs.

Since the core compiles into your module, a consumer does not even import fast_unionfind at runtime for these calls, only to construct the object that owns the arrays.

Thread safety

Concurrent finds are safe as long as no union is running at the same time. They do race, on path-compression stores, but harmlessly. Every such store points a node at one of its own ancestors, so the root stays reachable and unchanged. A lost write can only undo some compression, never send a walk somewhere wrong, and the aligned integer slots mean no store is ever half-written. This is what makes a lock-free parallel "are these already connected?" filter over a Kruskal edge list sound.

Unions must not run concurrently with each other or with finds. Serialise them.

The extension is declared freethreading_compatible, and nothing in it needs the GIL.

Performance

Two workloads over m = 4n random edges, written as Cython loops. The first, unite_all, is Kruskal's scan. The second, filter_scan, unites half the edges and then asks "already connected?" of all of them, which is two finds per edge. The comparison is the same algorithm written the usual way, as methods on a cdef class cimported from another extension module, where Cython must call through the imported vtable and nothing can inline.

Best of 50 runs at n ≤ 300 000 and best of 5 above that, on x86-64 with GCC -O3, Python 3.14 and Cython 3.3. The gain narrows at 10 million elements, where the 40 MB parent array no longer fits in cache and memory traffic dominates the call cost.

From Python, method calls cost about 42 ns each, which is almost all interpreter overhead and the same as calling a class method directly.

The class-based column was measured before that implementation was retired, so it cannot be re-run from this repository. python benchmarks/bench.py times this package's workloads and Python-level calls, compiling the workloads on first use.

License

BSD 3-clause; see LICENSE.