ABC218 E — Destruction

1# -*- coding: utf-8 -*-2 3 4from sys import int_info5 6 7class UnionFind:8    '''Represents a data structure that tracks a set of elements partitioned9       into a number of disjoint (non-overlapping) subsets.10    Landau notation: O(α(n)), where α(n) is the inverse Ackermann function.11    See:12    https://www.youtube.com/watch?v=zV3Ul2pA2Fw13    https://en.wikipedia.org/wiki/Disjoint-set_data_structure14    https://atcoder.jp/contests/abc120/submissions/444494215    '''16 17    def __init__(self, number_count: int):18        '''19        Args:20            number_count: The size of elements (greater than 2).21        '''22        self.parent_numbers = [-1 for _ in range(number_count)]23 24    def find_root(self, number: int) -> int:25        '''Follows the chain of parent pointers from number up the tree until26           it reaches a root element, whose parent is itself.27        Args:28            number: The trees id (0-index).29        Returns:30            The index of a root element.31        '''32        if self.parent_numbers[number] < 0:33            return number34 35        self.parent_numbers[number] = self.find_root(self.parent_numbers[number])36        return self.parent_numbers[number]37 38    def get_group_size(self, number: int) -> int:39        '''40        Args:41            number: The trees id (0-index).42        Returns:43            The size of group.44        '''45        return -self.parent_numbers[self.find_root(number)]46 47    def is_same_group(self, number_x: int, number_y: int) -> bool:48        '''Represents the roots of tree number_x and number_y are in the same49           group.50        Args:51            number_x: The trees x (0-index).52            number_y: The trees y (0-index).53        '''54        return self.find_root(number_x) == self.find_root(number_y)55 56    def merge_if_needs(self, number_x: int, number_y: int) -> bool:57        '''Uses find_root to determine the roots of the tree number_x and58           number_y belong to. If the roots are distinct, the trees are combined59           by attaching the roots of one to the root of the other.60        Args:61            number_x: The trees x (0-index).62            number_y: The trees y (0-index).63        '''64        x = self.find_root(number_x)65        y = self.find_root(number_y)66 67        if x == y:68            return False69 70        if self.get_group_size(x) >= self.get_group_size(y):71            self.parent_numbers[x] += self.parent_numbers[y]72            self.parent_numbers[y] = x73        else:74            self.parent_numbers[y] += self.parent_numbers[x]75            self.parent_numbers[x] = y76        return True77 78 79def main():80    import sys81 82    input = sys.stdin.readline83 84    n, m = map(int, input().split())85    edges = list()86    uf = UnionFind(n)87 88    for i in range(m):89        ai, bi, ci = map(int, input().split())90        edges.append((ci, ai - 1, bi - 1))91    92    ans = 093 94    for ci, ai, bi in sorted(edges):95        if not uf.is_same_group(ai, bi):96            uf.merge_if_needs(ai, bi)97        else:98            if ci > 0:99                ans += ci100    101    print(ans)102 103 104if __name__ == "__main__":105    main()106