ABC399 C — Make it Forest

1# -*- coding: utf-8 -*-2 3 4from typing import List5 6 7class UnionFind:8    """Represents a data structure that tracks a set of elements partitioned9       into a number of disjoint (non-overlapping) subsets.10 11    Landau notation: O(α(n)), where α(n) is the inverse Ackermann function.12 13    See:14    https://www.youtube.com/watch?v=zV3Ul2pA2Fw15    https://en.wikipedia.org/wiki/Disjoint-set_data_structure16    https://atcoder.jp/contests/abc120/submissions/444494217    https://atcoder.jp/contests/abc292/submissions/3941007518    https://github.com/not522/ac-library-python/blob/master/atcoder/dsu.py19    """20 21    def __init__(self, number_count: int) -> None:22        """23        Args:24            number_count: The size of elements (greater than 2).25        """26        self.number_count = number_count27        self.parent_numbers = [-1 for _ in range(number_count)]28        self.edge_count = [0 for _ in range(number_count)]29        self.group_count = number_count30 31    def find_root(self, number: int) -> int:32        """Follows the chain of parent pointers from number up the tree until33           it reaches a root element, whose parent is itself.34        Args:35            number: The trees id (0-index).36 37        Returns:38            The index of a root element.39        """40        if self.parent_numbers[number] < 0:41            return number42 43        self.parent_numbers[number] = self.find_root(self.parent_numbers[number])44        return self.parent_numbers[number]45 46    def get_group_size(self, number: int) -> int:47        """48        Args:49            number: The trees id (0-index).50 51        Returns:52            The size of group.53        """54        return -self.parent_numbers[self.find_root(number)]55 56    def is_same_group(self, number_x: int, number_y: int) -> bool:57        """Represents the roots of tree number_x and number_y are in the same58           group.59        Args:60            number_x: The trees x (0-index).61            number_y: The trees y (0-index).62        """63        return self.find_root(number_x) == self.find_root(number_y)64 65    def merge_if_needs(self, number_x: int, number_y: int) -> bool:66        """Uses find_root to determine the roots of the tree number_x and67           number_y belong to. If the roots are distinct, the trees are combined68           by attaching the roots of one to the root of the other.69        Args:70            number_x: The trees x (0-index).71            number_y: The trees y (0-index).72        """73        x = self.find_root(number_x)74        y = self.find_root(number_y)75 76        self.edge_count[x] += 177 78        if x == y:79            return False80 81        self.group_count -= 182 83        if self.parent_numbers[x] > self.parent_numbers[y]:84            x, y = y, x85 86        self.parent_numbers[x] += self.parent_numbers[y]87        self.parent_numbers[y] = x88        self.edge_count[x] += self.edge_count[y]89        return True90 91    def get_roots(self) -> List[int]:92        return [i for i, x in enumerate(self.parent_numbers) if x < 0]93 94    def get_groups(self) -> List[List[int]]:95        roots: List[int] = [self.find_root(i) for i in range(self.number_count)]96        groups: List[List[int]] = [[] for _ in range(self.number_count)]97 98        for i in range(self.number_count):99            groups[roots[i]].append(i)100 101        return list(filter(lambda g: g, groups))102 103    def get_edge_count(self, number: int) -> int:104        return self.edge_count[number]105 106    def get_group_count(self) -> int:107        return self.group_count108 109 110def main():111    import sys112 113    input = sys.stdin.readline114 115    n, m = map(int, input().split())116    uf = UnionFind(n)117    ans = 0118 119    for _ in range(m):120        ui, vi = map(int, input().split())121        ui -= 1122        vi -= 1123 124        if uf.is_same_group(ui, vi):125            ans += 1126        else:127            uf.merge_if_needs(ui, vi)128 129    print(ans)130 131 132if __name__ == "__main__":133    main()134