# Weighted union-find with path compression in Python 3

I implemented the Union-Find data structure and would love to receive some feedback about my code. I have just begun programming and welcome any feedback or constructive criticism which would improve my code quality and professionalism.

class UnionFind:

def __init__(self, N=1001):
self.N = N
self.id = [i for i in range(self.N)]
self._rank = [1] * self.N

def __root(self, x):
while x != self.id[self.id[x]]:
self.id[x] = self.id[self.id[x]]
x = self.id[x]
return x

def find(self, p, q):
return self.__root(p) == self.__root(q)

def unite(self, p, q):
i = self.__root(p)
j = self.__root(q)
if self._rank[j] > self._rank[i]:
i, j = j, i
self.id[j] = i
self._rank[i] += self._rank[j]


I was working on the problem of finding cycles in undirected graphs and realised union-find could be used to efficiently solve this problem and hence, implemented it.

In this instance the input would be a list of [u, v] pairs where u and v are integers and u < v. u and v are the node numbers of the graph under consideration.

For example the input [[1, 3], [2, 3], [1, 2]] would represent

Finding cycles

The idea is as follows

iterate over the [u, v] pairs
if find(u, v): there exists a cycle since [u, v] belong to the same connected component
else: unite(u, v)

For the example graph [[1, 3], [2, 3], [1, 2]] above, the code would return on [1, 2] since adding this would result in a cycle.

• Why does find in wikipedia have one argument but your implementation has two? "Find(x) follows the chain of parent pointers from x upwards through the tree until an element is reached whose parent is itself. This element is the root of the tree and is the representative member of the set to which x belongs, and may be x itself." What is the objective of this change? Commented Oct 3, 2017 at 20:40
• Thank you for pointing this out. As I have mentioned in the question, I wrote this code with the goal of finding cycles in the graph and hence the find took two parameters u,v where u and v represented the node numbers of the graph. Since this was problem specific I used the Wiki implementation as a reference, but now I have changed the signature so that the code sticks to the original algorithm. Do have a look at the new version of the code. Thanks. Commented Oct 5, 2017 at 20:04
• Please don't change the code after receiving an answer; see faq. Commented Oct 6, 2017 at 5:32
• Oh I'm sorry I was not aware of this. I'll keep this in mind moving forward. Commented Oct 6, 2017 at 6:49

## 2 Answers

• unite should check if p and q are in the same set already and do nothing in that case.
• Using dictionaries instead of lists as internal data structures would make this more flexible. Now the elements must be from range(N) and N has to be known at the start. Wikipedia suggests a MakeSet function which you could implement like this, and your other functions would continue to work unchanged:

def __init__(self):
self.id = {}
self._rank = {}

def make_set(self, element):
self.id[element] = element
self._rank[element] = 1

• id is a non-obvious name for a data structure stores the parent of each node in a tree. Why not parent?

• A single leading underscore is conventionally used to indicate private data.
• It is possible condense these two lines to one and avoid the extra lookup of self.id[x]

self.id[x] = self.id[self.id[x]]
x = self.id[x]


with a chained assignment. However, you risk confusing a reader who is not well familiar with the order of evaluation. The right hand side is evaluated first, and the assignments are carried out starting from self.id[x] on the left, and the value of x changes after:

self.id[x] = x = self.id[self.id[x]]

• I would avoid condensing the statements because it can be confusing because x appears in two parts. The OP way is simpler. The other points are good Commented Oct 3, 2017 at 20:38
• @Caridorc Edited to point out the possible confusion. IMO the shorter version is easier to read, as long as you know how it works. Commented Oct 4, 2017 at 5:02
• @JanneKarila Thanks a lot for your comments. I have taken into consideration a few of your suggestions like refactoring unite function, renaming id as parent. I was not clear what you meant when you said when you suggested to use dictionary as the internal datastructure. Could you please elaborate on this? Also, regarding the assignment I feel self.id[x] = x = self.id[self.id[x]] is less readable than self.id[x] = self.id[self.id[x]] x = self.id[x] (Of course YMMV) Commented Oct 5, 2017 at 20:14

Thank you Janne Karila and Caridorc for your suggestions.I have incorporated some of your changes into the code and this is the revised code. I welcome any further suggestions or improvements. Thank you for your time.

class UnionFind:

def __init__(self, N=1001):
self.N = N
self.parent = [i for i in range(self.N)]
self._rank = [1] * self.N

def find(self, x):
while x != self.parent[self.parent[x]]:
self.parent[x] = self.parent[self.parent[x]]
x = self.parent[x]
return x

def unite(self, p, q):
i = self.find(p)
j = self.find(q)
if i == j:
return
if self._rank[j] > self._rank[i]:
i, j = j, i
self.parent[j] = i
self._rank[i] += self._rank[j]

• In the unite function I'd use if i != j: .... It would streamline the logic and remove extra return statement. Commented Sep 12, 2018 at 6:15