Weighted union-find with path compression in Python 3

Question

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.

Answer 1

• 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]]
  

Answer 2

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]