Union Find em python + descrição em PT-BR

union-find.py

# n : numero de nos
# q : numero de operacoes
#     operacoes:
#         1, u, v : conecta conjuntos de u e v
#         2, u    : impreme tamanho do conjunto de u
#         3, u    : imprime par de u
#
n, q = map(int, raw_input().split())

# -------------------- UNION FIND --------------------------
# para cada no i de 1 a n, o par[i] = i
par = range(n+1)
# rank inicialmente eh 1, e tamanho de cada conjunto (s) tbm eh 1
rank = [1] * (n+1)
s = [1] * (n+1) 

# funcao que retorna o par de um no, eh a funcao mais importante!!!
def get_par(x):
	# se x eh par[x] entao achamos o par de todo um conjunto
	if x == par[x]:
		return par[x]
	else:
		# caso contrario usamos path compression para deixar mais eficiente
		# assim atualizamos o par do no e o retornamos
		par[x] = get_par(par[x])
		return par[x]

# dois nos estao no mesmo conjunto se tem o mesmo par
# perceba que usamos a funcao get_par para verificar, essa funcao utilizara
# path compression
def is_same_set(x, y):
	return get_par(x) == get_par(y)

# conecta dois conjuntos
def connect(x, y):
	# se ainda nao estiverem conectados...
	if not is_same_set(x, y):
		par_x = get_par(x)
		par_y = get_par(y)
		# o de maior rank sera escolhido para "integrar" em si o outro conjunto
		# se tiver empate adicionamos 1 ao rank de algum deles (podemos escolher)
		# e tambem atualizamos o tamanho do conjunto.
		if rank[par_x] > rank[par_y]:
			par[par_y] = par_x 
			s[par_x] += s[par_y]
		elif rank[par_y] > rank[par_x]:
			par[par_x] = par_y
			s[par_y] += s[par_x]
		else:
			rank[par_x] += 1
			s[par_x] += s[par_y]
			par[par_y] = par_x

# pega o tamanho do conjunto (SEMPRE EM RELACAO AO PAR!!!)
def get_size(x):
	return s[get_par(x)]

for i in xrange(q):
	query = map(int, raw_input().split())
	if query[0] == 1:
		print 'conectando', query[1], query[2]
		connect(query[1], query[2])
	elif query[0] == 2:
		print 'tamanho do conjunto que', query[1], 'pertence:',
		print get_size(query[1])
	else:
		print 'par de', query[1], 'eh:',
		print get_par(query[1])