Skip to content

Instantly share code, notes, and snippets.

@shuuchen
Last active June 20, 2019 14:31
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save shuuchen/6abec81d23d7e04312ab4b6e9c44d79c to your computer and use it in GitHub Desktop.
Save shuuchen/6abec81d23d7e04312ab4b6e9c44d79c to your computer and use it in GitHub Desktop.
An implementation of union find tree
class UT:
def __init__(self, n):
self.ps = list(range(n))
self.rs = [0] * n
def find(self, x):
if self.ps[x] != x:
self.ps[x] = self.find(self.ps[x])
return self.ps[x]
def union(self, x, y):
px, py = self.find(x), self.find(y)
if px == py:
return
rx, ry = self.rs[px], self.rs[py]
if rx > ry:
self.ps[py] = px
elif rx < ry:
self.ps[px] = py
else:
self.ps[px] = py
self.rs[py] += 1
def get(self, n):
return sum(1 for i, p in enumerate(self.ps[:n]) if p == i)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment