ABC328 F — Good Set Query

1# -*- coding: utf-8 -*-2 3 4class WeightedUnionFind:5    """Represents a data structure that tracks a set of elements partitioned6       into a number of disjoint (non-overlapping) subsets.7    Landau notation: O(α(n)), where α(n) is the inverse Ackermann function.8 9    See:10    https://www.youtube.com/watch?v=zV3Ul2pA2Fw11    https://en.wikipedia.org/wiki/Disjoint-set_data_structure12    https://atcoder.jp/contests/abc120/submissions/444494213    https://qiita.com/drken/items/cce6fc5c579051e64fab14    """15 16    def __init__(self, number_count: int):17        """18        Args:19            number_count: The size of elements (greater than 2).20        """21        self.parent_numbers = [i for i in range(number_count)]22        self.rank = [0 for _ in range(number_count)]23        self.diff_weight = [0 for _ in range(number_count)]24 25    def find_root(self, number: int) -> int:26        """Follows the chain of parent pointers from number up the tree until27           it reaches a root element, whose parent is itself.28        Args:29            number: The trees id (0-index).30        Returns:31            The index of a root element.32        """33        if self.parent_numbers[number] == number:34            return number35        else:36            parent_number = self.parent_numbers[number]37            root = self.find_root(parent_number)38            self.diff_weight[number] += self.diff_weight[parent_number]39            self.parent_numbers[number] = root40            return root41 42    def calc_weight(self, number: int) -> int:43        """Calculate the weight of the node.44        Args:45            number: The trees id (0-index).46        Returns:47            The weight of the node.48        """49        self.find_root(number)50        return self.diff_weight[number]51 52    def is_same_group(self, number_x: int, number_y: int) -> bool:53        """Represents the roots of tree number_x and number_y are in the same54           group.55        Args:56            number_x: The trees x (0-index).57            number_y: The trees y (0-index).58        """59        return self.find_root(number_x) == self.find_root(number_y)60 61    def merge_if_needs(self, number_x: int, number_y: int, weight: int) -> bool:62        """Uses find_root to determine the roots of the tree number_x and63           number_y belong to. If the roots are distinct, the trees are64           combined by attaching the roots of one to the root of the other.65        Args:66            number_x: The trees x (0-index).67            number_y: The trees y (0-index).68            weight  : Cost between nodes.69        """70        # Correct the difference between the weight of root and number_x,71        # number_y.72        weight += self.calc_weight(number_x)73        weight -= self.calc_weight(number_y)74        root_x, root_y = self.find_root(number_x), self.find_root(number_y)75 76        if root_x == root_y:77            return False78 79        if self.rank[root_x] < self.rank[root_y]:80            root_x, root_y = root_y, root_x81            weight = -weight82        if self.rank[root_x] == self.rank[root_y]:83            self.rank[root_x] += 184 85        self.parent_numbers[root_y] = root_x86        self.diff_weight[root_y] = weight87        return True88 89    def calc_cost(self, from_x: int, to_y: int) -> int:90        """Calculate cost between nodes.91        Args:92            from_x: The trees x (0-index).93            to_y  : The trees y (0-index).94        Returns:95            Cost between nodes.96        """97        return self.calc_weight(to_y) - self.calc_weight(from_x)98 99 100def main():101    import sys102 103    input = sys.stdin.readline104 105    n, q = map(int, input().split())106    uf = WeightedUnionFind(n)107    ans = list()108 109    for i in range(q):110        ai, bi, di = map(int, input().split())111        ai -= 1112        bi -= 1113 114        if not uf.is_same_group(ai, bi):115            uf.merge_if_needs(ai, bi, di)116 117        cost = uf.calc_cost(ai, bi)118 119        if cost == di:120            ans.append(i + 1)121 122    print(*ans)123 124 125if __name__ == "__main__":126    main()127