# Union-find disjoint set solution to calculating sizes of connected components

I'm given a distance and a list of points in a plane, where I have to return a sorted list of the size of each connected component where each edge is shorter or equal to the distance threshold.

I've implemented a union-find disjoint set as a solution:

import array
from collections import Counter

class DisjointSet:
def __init__(self, n):
self.n = n
self.parent = array.array('I', range(self.n))
self.rank = array.array('I', [0] * self.n)

def sizes_of_unique_sets(self):
count = Counter(self.find(i) for i in range(self.n))
return list(count.values())

def find(self, a):
if self.parent[a] != a:
self.parent[a] = self.find(self.parent[a])
return self.parent[a]

def union(self, a, b):
a_root = self.find(a)
b_root = self.find(b)
if a_root == b_root:
return

if self.rank[a_root] > self.rank[b_root]:
self.parent[b_root] = a_root
elif self.rank[a_root] < self.rank[b_root]:
self.parent[a_root] = b_root
else:
self.parent[b_root] = a_root
self.rank[a_root] += 1


To compute the solution, I iterate the list of points twice (the second loop starts from the first index + 1)

def disjoint_set_compute(distance, points):
N = len(points)
disjoint_set = DisjointSet(N)

for i in range(N - 1):
for j in range(i + 1, N):
if points[i].inside_ball(points[j], distance) and disjoint_set.find(i) != disjoint_set.find(j):
disjoint_set.union(i, j)

res = sorted(disjoint_set.sizes_of_unique_sets(), reverse=True)
return res


When benchmarked, I get 5 seconds from 1,000 points on 6 virtual cores, which is kind of sad, but it's even worse that I can't find a way to further optimise it, other than using a spatial data structure like a Quadtree or KD-tree, which I will work on further. I can also try a BFS approach.

Is there any way to improve this current implementation?

# lint

    N = len(points)


Pep-8 asks that you assign to downcased n, please.

# verb

                disjoint_set.union(i, j)


Hmmm, that's slightly odd. It's clear we're evaluating for side effects. It doesn't seem to fit in with the familiar union operator in standard library, which produces a new result object.

Consider a more explicit name that is clearly a verb, perhaps .add_to_union.

# docstrings

The DisjointSet class really deserves a """docstring""", even if it's just a wiki ref. And the parent and rank datastructures need a little love, in the class docstring or even just as a # comment.

It's clear what sizes_of_unique_sets does, but a docstring really ought to explain it so caller knows what return value to expect without examining the code.

It turns out that find and union deserve docstrings, too.

There's probably something you'd like to say about big-Oh time complexity of selected operations, such as union.

# recursion

        if self.parent[a] != a:
self.parent[a] = self.find(self.parent[a])


In many ways this is beautiful.

But consider either

1. coding this in iterative fashion, or
2. documenting how n and sys.getrecursionlimit() are related

Even if this was implemented in Scheme there's no opportunity for TCO. Large datastructures will trigger RecursionError, essentially a stack overflow.

I am perfectly fine with the fact that .find has side effects. But that should be mentioned in the docstring, as simple inspection of the signature would suggest it is readonly.

# invariants

What is always true about parent ? And about rank ? Write them down.

In the course of writing them down, you might possibly find that the definition is inconvenient, and you prefer to do all the mutating within .union, rather than handing some of the responsibility to .find.

Do cite the references that inspired your choice of datastructure and of algorithm.

# speed

5 seconds from 1,000 points on 6 virtual cores

It's unclear why VCPU > 1 would be relevant, since there's just one GIL.

Consider coding some of this library in a language like Rust if you want .union to go faster. Though even in that case, it would still seem like there's very little opportunity to spin up worker threads and get anything useful done with them, unlike e.g. numpy finding the inverse of a large matrix.

Adding numba @jit decorators would be pretty low effort. You have already adopted a sensible array datastructure (rather than the cache misses of chasing list object pointers), so that's good.

Do consider sending your (x, y) coords to a relational database that has already implemented spatial indexes for you. I have had excellent results with radius-based queries against PostGIS via sqlalchemy. For large number of points, using appropriate datastructure really makes a huge difference. There's just no comparison between having pair of x- and y-indexes versus a combined spatial xy-index.

I love sqlite but have never used its spatial R-tree, nor the SpatiaLite extension. It would be thirty minutes well invested to try them out.

This class achieves many of its design goals. Some documentation deficiencies really ought to be addressed prior to merging down to main. Given that, I would be happy to delegate or accept maintenance tasks on this codebase.

I get 5 seconds from 1,000 points

There are about n²/2 inside_ball queries and possibly a similar number of union/find operations. So, 10 microseconds per pair, on average. For a language that isn't Python that would be a lot, but for Python it does not sound too outrageous to me.

disjoint_set.find(i) != disjoint_set.find(j)

Is redundant, union does that anyway. But it should not have a large impact. Calling find here makes, thanks to path compression, the find inside union cheap. Either way the real work is done somewhere, but only once, and the only thing that may change is the amount of overhead.
A DisjointSet based on a couple of arrays of integers is good, perhaps some details could be slightly improved.
find does not need to be recursive (which walks the path twice, once up and then back down), that is an obvious way to do it, but there are alternatives. The "classic" one is still a two-pass algorithm: find the ultimate parent (iteratively), then walk back up and update everyones parent. As noted on wikipedia, there are one-pass iterative alternatives, path splitting and path halving. I don't know how good they are, but you can try them. They walk the path only once, but the trade-off is that they do less path compression.