ABC380Eを解いた【UnionFind】

筆者はレート800前後の茶～緑コーダ

ABC380のE問題を解いていく

実装コード

UnionFindを使う
集団に左右端の位置を覚えさせておき
塗り替えた場所の隣と色が同じになった場合は
それらの集合と合体させればOK

main.py

import sys
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rI1(): return (int(sys.stdin.readline().rstrip())-1)
def rLI1(): return list(map(lambda a:int(a)-1,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
def err(*args): print(*args, file=sys.stderr)


class UnionFind:
    def __init__(self, n):
        self.parent = [-1] * -~n
        self.rank = [1] * -~n
        self.label = list(range(-~n))
        self.left = list(range(-~n))
        self.right = list(range(-~n))
        self.count = n
        self.c_cnt = [1 for i in range(-~n)]

    def root(self, x):
        while (p := self.parent[x]) != -1:
            x = p
        return x

    def union(self, x, y):
        xr = self.root(x)
        yr = self.root(y)
        nl = min(self.left[xr],self.left[yr])
        nr = max(self.right[xr],self.right[yr])
        if xr == yr:
            return
        elif (rx := self.rank[xr]) < (ry := self.rank[yr]):
            self.parent[xr] = yr
            self.rank[yr] += rx
            self.left[yr] = nl
            self.right[yr] = nr
            # err(self.label[yr],nl,nr)
        else:
            self.parent[yr] = xr
            self.rank[xr] += ry
            self.left[xr] = nl
            self.right[xr] = nr
            # err(self.label[xr],nl,nr)
        self.count -= 1

    def same(self, x, y):
        return self.root(x) == self.root(y)
    
    def get_left(self, x):
        return self.left[self.root(x)]
    def get_right(self, x):
        return self.right[self.root(x)]
    def get_label(self, x):
        return self.label[self.root(x)]
    def set_label(self, x, c):
        xr = self.root(x)
        rx = self.rank[xr]
        xc = self.get_label(x)
        self.label[xr] = c
        self.c_cnt[xc] -= rx
        self.c_cnt[c] += rx
    


def main():
    N, Q = rLI()
    
    UF= UnionFind(N)
    
    for _ in range(Q):
        q = rLI()
        if q[0] == 1:
            x,c=q[1:]
            UF.set_label(x,c)
            lx = UF.get_left(x)
            rx = UF.get_right(x)
            if lx>1 and UF.get_label(lx-1) == c:
                UF.union(x,lx-1)
            if rx<N and UF.get_label(rx+1) == c:
                UF.union(x,rx+1)
            # err([UF.get_label(i+1) for i in range(N)])
        if q[0] == 2:
            c=q[1]
            print(UF.c_cnt[c])
            
    
if __name__ == '__main__':
    main()

まとめ

ABC380Eを解いた。実装にはUnionFindを使用した。

確か時間内に解けなかったのが残念だった記憶だが、
実装自体にはなんら苦労はなかったので
もっと早く実装する能力が求められるようだ。