ABC380 E — 1D Bucket Tool

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 get_group_size(self, number: int) -> int:46        """47        Args:48            number: The trees id (0-index).49 50        Returns:51            The size of group.52        """53        return -self.parent_numbers[self.find_root(number)]54 55    def is_same_group(self, number_x: int, number_y: int) -> bool:56        """Represents the roots of tree number_x and number_y are in the same57           group.58        Args:59            number_x: The trees x (0-index).60            number_y: The trees y (0-index).61        """62        return self.find_root(number_x) == self.find_root(number_y)63 64    def merge_if_needs(self, number_x: int, number_y: int) -> bool:65        """Uses find_root to determine the roots of the tree number_x and66           number_y belong to. If the roots are distinct, the trees are combined67           by attaching the roots of one to the root of the other.68        Args:69            number_x: The trees x (0-index).70            number_y: The trees y (0-index).71        """72        x = self.find_root(number_x)73        y = self.find_root(number_y)74 75        self.edge_count[x] += 176 77        if x == y:78            return False79 80        self.group_count -= 181 82        if self.parent_numbers[x] > self.parent_numbers[y]:83            x, y = y, x84 85        self.parent_numbers[x] += self.parent_numbers[y]86        self.parent_numbers[y] = x87        self.edge_count[x] += self.edge_count[y]88        return True89 90    def get_roots(self) -> List[int]:91        return [i for i, x in enumerate(self.parent_numbers) if x < 0]92 93 94def main():95    import sys96 97    input = sys.stdin.readline98 99    n, q = map(int, input().split())100    n += 2  # 番兵(0番目、n + 1番目)101    uf = UnionFind(n)102    # 各マスの配色とその個数、グループの左右端を管理103    counts = [1] * n104    colors = [i for i in range(n)]105    left = [i for i in range(n)]106    right = [i for i in range(n)]107 108    for _ in range(q):109        qi = list(map(int, input().split()))110        type = qi[0]111 112        if type == 1:113            x, c = qi[1], qi[2]114 115            x = uf.find_root(x)116            size = uf.get_group_size(x)117 118            # マスの色と個数を更新119            counts[colors[x]] -= size120            colors[x] = c121            counts[colors[x]] += size122 123            # 隣接グループとのマージができるか判定124            # 左側125            li = uf.find_root(left[x] - 1)126 127            if colors[li] == c:128                nl, nr = left[li], right[x]  # 一時的に別変数に保存129                uf.merge_if_needs(li, x)130                x = uf.find_root(x)131                left[x], right[x], colors[x] = nl, nr, c132 133            # 右側134            ri = uf.find_root(right[x] + 1)135 136            if colors[ri] == c:137                nl, nr = left[x], right[ri]  # 一時的に別変数に保存138                uf.merge_if_needs(ri, x)139                x = uf.find_root(x)140                left[x], right[x], colors[x] = nl, nr, c141        else:142            c = qi[1]143            print(counts[c])144 145 146if __name__ == "__main__":147    main()148