from array import array

import os

class DisjointSets(object):

def __init__(self, size):

self.hierarchy = array('l', [0 for _ in range(size)])

self.names = {}

def add(self, item_id, name):

self.hierarchy[item_id] = -1

self.names[item_id] = name

def union(self, root1, root2):

if self.hierarchy[root1] <= self.hierarchy[root2]: # root1 is a larger tree, since roots are -size of tree

self.hierarchy[root1] += self.hierarchy[root2] # adding to increase negative value

self.hierarchy[root2] = root1

else:

self.hierarchy[root2] += self.hierarchy[root1]

self.hierarchy[root1] = root2

def find(self, item_id):

'''

Finds the root of the set of which item_id is a member.

To speed up subsequent finds, updates the parent to root at each level in path to root.

:param item_id:

:return: integer representative of set

'''

if self.hierarchy[item_id] < 0:

return item_id

else:

# path compression

self.hierarchy[item_id] = self.find(self.hierarchy[item_id])

return self.hierarchy[item_id]

def __str__(self):

ret_str = repr(self.hierarchy)

ret_str += os.linesep + repr(self.names)

return ret_str

def main():

ds = DisjointSets(10)

ds.add(0, "Microsoft")

ds.add(1, "WebTV")

ds.add(2, "Google")

ds.add(3, "DeepMind")

ds.add(4, "Skype")

ds.add(5, "Uber")

ds.add(6, "TinyCo")

ds.add(7, "TeenyTinyCo")

ds.union(0, 1)

ds.union(2, 3)

ds.union(0, 4)

ds.union(4, 6)

ds.union(6, 7)

# print(ds)

assert(ds.find(3) == 2)

# print(ds)

if __name__ == '__main__':

main()