ABC378 F — Add One Edge 2

1# -*- coding: utf-8 -*-2 3from typing import List4 5 6class UnionFind:7    """Represents a data structure that tracks a set of elements partitioned8       into a number of disjoint (non-overlapping) subsets.9 10    Landau notation: O(α(n)), where α(n) is the inverse Ackermann function.11 12    See:13    https://www.youtube.com/watch?v=zV3Ul2pA2Fw14    https://en.wikipedia.org/wiki/Disjoint-set_data_structure15    https://atcoder.jp/contests/abc120/submissions/444494216    https://atcoder.jp/contests/abc292/submissions/3941007517    https://github.com/not522/ac-library-python/blob/master/atcoder/dsu.py18    """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.number_count = number_count26        self.parent_numbers = [-1 for _ in range(number_count)]27        self.edge_count = [0 for _ in range(number_count)]28        self.group_count = number_count29 30    def find_root(self, number: int) -> int:31        """Follows the chain of parent pointers from number up the tree until32           it reaches a root element, whose parent is itself.33        Args:34            number: The trees id (0-index).35 36        Returns:37            The index of a root element.38        """39        if self.parent_numbers[number] < 0:40            return number41 42        self.parent_numbers[number] = self.find_root(self.parent_numbers[number])43        return self.parent_numbers[number]44 45    def merge_if_needs(self, number_x: int, number_y: int) -> bool:46        """Uses find_root to determine the roots of the tree number_x and47           number_y belong to. If the roots are distinct, the trees are combined48           by attaching the roots of one to the root of the other.49        Args:50            number_x: The trees x (0-index).51            number_y: The trees y (0-index).52        """53        x = self.find_root(number_x)54        y = self.find_root(number_y)55 56        self.edge_count[x] += 157 58        if x == y:59            return False60 61        self.group_count -= 162 63        if self.parent_numbers[x] > self.parent_numbers[y]:64            x, y = y, x65 66        self.parent_numbers[x] += self.parent_numbers[y]67        self.parent_numbers[y] = x68        self.edge_count[x] += self.edge_count[y]69        return True70 71    def get_groups(self) -> List[List[int]]:72        roots: List[int] = [self.find_root(i) for i in range(self.number_count)]73        groups: List[List[int]] = [[] for _ in range(self.number_count)]74 75        for i in range(self.number_count):76            groups[roots[i]].append(i)77 78        return list(filter(lambda g: g, groups))79 80 81def main():82    import sys83 84    input = sys.stdin.readline85 86    n = int(input())87    degrees = [0] * n88    uv = list()89 90    for _ in range(n - 1):91        ai, bi = map(int, input().split())92        ai -= 193        bi -= 194        uv.append((ai, bi))95 96        degrees[ai] += 197        degrees[bi] += 198 99    # 次数3の頂点のうち、隣接しているものを同じ連結成分にまとめる100    uf = UnionFind(n)101    # 次数3 - 次数2の辺の数 (次数3の頂点に隣接している次数2の頂点を走査するため、添え字は次数3の頂点番号)102    counts = [0] * n103 104    for ui, vi in uv:105        if degrees[ui] == 3 and degrees[vi] == 3:106            uf.merge_if_needs(ui, vi)107        elif degrees[ui] == 3 and degrees[vi] == 2:108            counts[ui] += 1109        elif degrees[ui] == 2 and degrees[vi] == 3:110            counts[vi] += 1111 112    ans = 0113 114    # 同じ連結成分を対象として、次数3の頂点に隣接している次数2の頂点のペアを求める115    for group in uf.get_groups():116        count = 0117 118        for vertex in group:119            count += counts[vertex]120 121        ans += count * (count - 1) // 2122 123    print(ans)124 125 126if __name__ == "__main__":127    main()128