ABC372Eを解いた【UnionFind】

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

ABC372のE問題を解いていく

実装コード

kの最大値が10なのを利用して
UnionFindの各ノードにに長さ10のリストを乗せて
マージするたびにソートすれば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, k):
        self.parent = [-1] * -~n
        self.rank = [1] * -~n
        self.count = n
        self.k = k
        self.L = [[i]+[-1]*~-k 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)
        # err(xr,yr)
        if xr == yr: return
        XL = self.L[xr]
        YL = self.L[yr]
        ML = sorted(XL+YL,reverse=True)[:self.k]
        # err(ML)
        if (rx:=self.rank[xr]) < (ry:=self.rank[yr]):
            self.parent[xr] = yr
            self.rank[yr] += rx
            self.L[yr] = ML
        else:
            self.parent[yr] = xr
            self.rank[xr] += ry
            self.L[xr] = ML
        self.count -= 1
    def same(self,x,y):
        # err(self.root(x), self.root(y))
        return self.root(x) == self.root(y)
    
def main():
    N, Q = rLI()
    UF=UnionFind(N,10)
    for _ in range(Q):
        q, *I = rLI()
        if q==1:
            UF.union(*I)
        if q==2:
            v,k=I
            rv = UF.root(v)
            ans = UF.L[rv][k-1]
            print(ans)
    
    
if __name__ == '__main__':
    main()

まとめ

ABC372のE問題を解いた

当時はグラフの問題はできませんorzとか嘆いていた問題だけど、UnionFind一発で解けたので成長を感じた。