# DisjointSet class in Python

I would like to improve my DisjointSet class implementation in Python. I would be grateful for specific advice. I use it in online contests so it must be short, quick and reusable.

from copy import copy
class DisjointSet:

def __init__(self, s=set()):
self.__par = {x: x for x in s}
self.__csz = dict.fromkeys(s, 1)
self.__root = set(s)

def root(self, x=None):
if x is None:
return(copy(self.__root))
r = x
while r != self.__par[r]:
r = self.__par[r]
while x != r:
p = self.__par[x]
self.__par[x] = r
x = p
return(r)

if x not in self.__par:
self.__par[x] = x
self.__csz[x] = 1

def connected(self, x, y):
return(self.root(x) == self.root(y))

def connect(self, x, y):
rx, ry = self.root(x), self.root(y)
if rx == ry:
return(rx)
if self.__csz[rx] < self.__csz[ry]:
rx, ry = ry, rx
self.__par[ry] = rx
self.__csz[rx] += self.__csz[ry]
self.__root.remove(ry)
return(rx)

def size(self, x=None):
if x is None:
return(len(self.__root))
else:
return(self.__csz[self.root(x)])

def __len__(self):
return(len(self.__par))


Simple test

from random import randint
#ds = DisjointSet(xrange(1000000))
ds = DisjointSet()
for x in xrange(1000000):
print len(ds)
print ds.size()
for __ in xrange(1000000):
ds.connect(randint(0, 999999), randint(0, 999999))
print ds.root(0)
print sum([ds.size(x) for x in ds.root()])
print ds.size()

• Can you explain what the purpose of this code actually is? What it's supposed to do? Not just by providing test cases, but an actual, English-language explanation? Dec 8, 2017 at 15:27

I'll edit this once it's posted what the code is supposed to do.

The code is extremely cryptic. There are no comments; the variable & membervariable names are all abbreviations. Someone else looking at this code will have no idea what is supposed to be happening. At the very least there should be docstrings at the start of each function/method that describe what the function is used for and the purposes of each variable.

The only possible performance improvement I see is to combine the __par, __csz, and __root members since they all have the same keys. That way the same element doesn't need to be hashed multiple times. The data structure, instead of being split between three dictionaries each of which being a member variable/field, would just have one dictionary with appropriate keys and the values being a list of [parent, csz, root].

def __init__(self, s=set()):
self.__data = {x: [x, 1, x] for x in s}


If you need one of the members you just refer to self.__data[x][0] or whatever.