ABC333 D — Erase Leaves

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 = int(input())105    uf = UnionFind(n)106    neighbors = list()107 108    for _ in range(n - 1):109        ai, bi = map(int, input().split())110        ai -= 1111        bi -= 1112 113        if ai == 0:114            neighbors.append(bi)115        else:116            if not uf.is_same_group(ai, bi):117                uf.merge_if_needs(ai, bi)118 119    sizes = list()120 121    for neighbor in neighbors:122        root = uf.find_root(neighbor)123        sizes.append(uf.get_group_size(root))124 125    # print(sizes)126    ans = sum(sizes) - max(sizes) + 1127    print(ans)128 129 130if __name__ == "__main__":131    main()132