ABC392Eを解いた【UnionFind】

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

ABC392のE問題を解いていく

実装コード

UnionFindを使い、クラスタ同士もしくは自己ループをスペアとみなし、判別する。

また、つなげかえる回数はUnionFindのクラスタ数-1で、
スペアのケーブルをクラスタ別に分け、
多い順から、別のクラスタにつなげかえる、
もし使い切ったら、次に多いクラスタにあるケーブルを使えばOK

main.py

from sortedcontainers import SortedSet

from bisect import bisect_left, bisect_right, insort_left, insort_right
from collections import defaultdict, Counter, deque
from functools import reduce, lru_cache
from itertools import product, accumulate, groupby, combinations
import sys
import os
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())
IS_LOCAL = int(os.getenv("ATCODER", "0"))==0
err = (lambda *args, **kwargs: print(*args, **kwargs, file=sys.stderr)) if IS_LOCAL else (lambda *args, **kwargs: None)

class UnionFind:
    def __init__(self, n):
        self.parent = [-1] * -~n
        self.rank = [1] * -~n
        self.roots= set([-~i for i in range(n)])
        self.count = 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)
        if xr == yr: return
        elif (rx:=self.rank[xr]) < (ry:=self.rank[yr]):
            self.parent[xr] = yr
            self.rank[yr] += rx
            self.roots.remove(xr)
        else:
            self.parent[yr] = xr
            self.rank[xr] += ry
            self.roots.remove(yr)
        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, M = rLI()
    UF = UnionFind(N)
    spare = []
    for i0 in range(M):
        i = -~i0
        a,b = rLI()
        # スペアの判別
        if a!=b and UF.root(a) != UF.root(b):
            UF.union(a,b)
        else:
            spare.append((i,a,b))
    # つなげる個数
    print(UF.count-1)
    # err(*UF.roots)
    # err(*spare)
    if UF.count > 1:
        
        # スペアをクラスタ別に分ける
        r = defaultdict(deque)
        c = defaultdict(lambda: 0)
        for s in spare:
            i,u,_ = s
            ru = UF.root(u)
            r[ru].append(s)
            c[ru]+=1
        # err(dict(r))
        
        # 多いクラスタから使う
        K = list(sorted((k for k in UF.roots),reverse=True, key=lambda x: c[x]))
        x = K[0]
        cur = 0
        # err(K)
        
        # 付け替え
        for y in K[1:]:
            # err(UF.count,x,y)
            # err(UF.count,xr,yr)
            # 使い切ったらクラスタ変える
            if len(r[x]) == 0:
                cur += 1
                x = K[cur]
            i,a,b = r[x].popleft()
            print(i,b,y)
            
if __name__ == '__main__':
    main()

コメント

ケーブルの付け替えの実装で,TLE,REを多数だしたが、
余りが多いクラスタから付け替えるという発想は目からウロコだった