Spaces:
Runtime error
Runtime error
| from typing import List | |
| class UnionFind: | |
| def __init__(self, data: List, union_condition: callable): | |
| self.__data__ = data | |
| self.__union_condition__ = union_condition | |
| length = len(data) | |
| self.__parents__ = [i for i in range(length)] | |
| self.__ranks__ = [0] * length | |
| self.__unions__ = {} | |
| def __find_parent__(self, id: int): | |
| return self.__parents__[id] | |
| def __find_root__(self, id: int): | |
| parent = self.__find_parent__(id) | |
| while parent != id: | |
| id = parent | |
| parent = self.__find_parent__(id) | |
| return id | |
| def __union__(self, i: int, j: int): | |
| root_i = self.__find_root__(i) | |
| root_j = self.__find_root__(j) | |
| # if roots are different, let one be the parent of the other | |
| if root_i == root_j: | |
| return | |
| else: | |
| if self.__ranks__[root_i] <= self.__ranks__[root_j]: | |
| # root of i --> child | |
| self.__parents__[root_i] = root_j | |
| self.__ranks__[root_j] = max(self.__ranks__[root_j], self.__ranks__[root_i]+1) | |
| else: | |
| self.__parents__[root_j] = root_i | |
| self.__ranks__[root_i] = max(self.__ranks__[root_i], self.__ranks__[root_j]+1) | |
| def union_step(self): | |
| length = len(self.__data__) | |
| for i in range(length - 1): | |
| for j in range(i + 1, length): | |
| if self.__union_condition__(self.__data__[i], self.__data__[j]): | |
| self.__union__(i, j) | |
| for i in range(length): | |
| root = self.__find_root__(i) | |
| if root not in self.__unions__.keys(): | |
| self.__unions__[root] = [self.__data__[i]] | |
| else: | |
| self.__unions__[root].append(self.__data__[i]) | |
| def get_unions(self): | |
| return self.__unions__ | |