ABC350 D — New Friends

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    """19 20    def __init__(self, number_count: int) -> None:21        """22        Args:23            number_count: The size of elements (greater than 2).24        """25        self.parent_numbers = [-1 for _ in range(number_count)]26        self.edge_count = [0 for _ in range(number_count)]27        self.group_count = number_count28 29    def find_root(self, number: int) -> int:30        """Follows the chain of parent pointers from number up the tree until31           it reaches a root element, whose parent is itself.32        Args:33            number: The trees id (0-index).34 35        Returns:36            The index of a root element.37        """38        if self.parent_numbers[number] < 0:39            return number40 41        self.parent_numbers[number] = self.find_root(self.parent_numbers[number])42        return self.parent_numbers[number]43 44    def get_group_size(self, number: int) -> int:45        """46        Args:47            number: The trees id (0-index).48 49        Returns:50            The size of group.51        """52        return -self.parent_numbers[self.find_root(number)]53 54    def is_same_group(self, number_x: int, number_y: int) -> bool:55        """Represents the roots of tree number_x and number_y are in the same56           group.57        Args:58            number_x: The trees x (0-index).59            number_y: The trees y (0-index).60        """61        return self.find_root(number_x) == self.find_root(number_y)62 63    def merge_if_needs(self, number_x: int, number_y: int) -> bool:64        """Uses find_root to determine the roots of the tree number_x and65           number_y belong to. If the roots are distinct, the trees are combined66           by attaching the roots of one to the root of the other.67        Args:68            number_x: The trees x (0-index).69            number_y: The trees y (0-index).70        """71        x = self.find_root(number_x)72        y = self.find_root(number_y)73 74        self.edge_count[x] += 175 76        if x == y:77            return False78 79        self.group_count -= 180 81        if self.parent_numbers[x] > self.parent_numbers[y]:82            x, y = y, x83 84        self.parent_numbers[x] += self.parent_numbers[y]85        self.parent_numbers[y] = x86        self.edge_count[x] += self.edge_count[y]87        return True88 89    def get_roots(self) -> List[int]:90        return [i for i, x in enumerate(self.parent_numbers) if x < 0]91 92    def get_edge_count(self, number: int) -> int:93        return self.edge_count[number]94 95    def get_group_count(self) -> int:96        return self.group_count97 98 99def main():100    import sys101 102    input = sys.stdin.readline103 104    n, m = map(int, input().split())105    graph = [[] for _ in range(n)]106    uf = UnionFind(n)107 108    # 各連結成分に対して、取りうる辺の最大数 - 元々の辺の数109    for _ in range(m):110        ai, bi = map(int, input().split())111        ai -= 1112        bi -= 1113 114        graph[ai].append(bi)115        graph[bi].append(ai)116 117        if not uf.is_same_group(ai, bi):118            uf.merge_if_needs(ai, bi)119 120    ans = 0121 122    for root in uf.get_roots():123        size = uf.get_group_size(root)124        count = size * (size - 1) // 2  # nC2125        ans += count126 127    print(ans - m)128 129 130if __name__ == "__main__":131    main()132