ABC350備忘録【C~D】

はじめに

D問題のUnion-Findのコードは競プロの鉄則のgithubをもとに自作関数を含めて作成したと思うので、参考書籍として競技プログラミングの鉄則を置いておきます。

C - Sort

問題

考えたこと

(index, 数値)で状態を格納しておき、indexと数値が異なるならば同じようになるよう交換するという操作を繰り返すということを行いました。操作の概念図は図1となります。

1ループでは終わらない時があるので、変更毎の数列は保存しておき、下記のtargetと等しくなるまでループ操作を続行させます。

target = [i + 1  for i in range(n)]

ACコード

c.py

n = int(input())
A = list(map(int, input().split()))

ans = []
target = [i + 1 for i in range(n)]
while True:
    for i in range(n):
        if i + 1 != A[i]:
            ans.append([i + 1, A[i]])
            Ai = A[i]
            A[i] = A[Ai - 1]
            A[Ai - 1] = Ai
    if A == target:
        break

print(len(ans))
for a in ans:
    a.sort()
    print(*a)

D - New Friends

問題

補足

少しグラフの用語を使っているので、基礎知識を知りたい方はこちら

考えたこと

今回の例では、各連結成分に対して最終的なグラフは完全グラフとなります。よって、各連結成分において、入力時のグラフと完全グラフで作成された補グラフの辺の数が操作の回数となります。
試しに以下の入力例を考えてみましょう。
入力例：

5 4
1 2
2 3
3 4
4 5

入力例から補グラフまでのは図2のようになります。

図から最終的なグラフは完全グラフとなっていることがわかります。
ここでも、一つの連結成分に注目してみます。点の数を$x_i$としたとき、完全グラフの辺の数は$x_i + x_i角形の対角線$です。さらに、$x_i$角形の対角線は$\frac{x_i(x_i -3)}{2}$と表されることと、入力時グラフの辺の数より、

補グラフの辺の数 = x_i +  \frac{x_i(x_i -3)}{2} - 入力時グラフの辺の数

となります。これを各連結成分ごとに和をとれば答えとなります。
グラフの作成や辺の本数の数え上げについては、Union-Findを用いて実装しました。

ACコード

d.py

class UnionFind:
    def __init__(self, N: int) -> None:
        self.N = N
        self.par = [-1 for _ in range(N)]
        self.siz = [1 for _ in range(N)]

    def root(self, x: int) -> int:
        while True:
            if self.par[x] == -1:
                break
            x = self.par[x]
        return x

    def unite(self, u: int, v: int) -> None:
        RootU = self.root(u)
        RootV = self.root(v)
        if RootU != RootV:
            if self.siz[RootU] < self.siz[RootV]:
                self.par[RootU] = RootV
                self.siz[RootV] += self.siz[RootU]
            else:
                self.par[RootV] = RootU
                self.siz[RootU] += self.siz[RootV]

    def same(self, u: int, v: int) -> bool:
        # 同じ根かどうか
        return self.root(u) == self.root(v)

    def root_list(self) -> set:
        # 親のsetを取得
        return {self.root(i) for i in range(self.N)}

    def count_root_v(self) -> list:
        # 各根に対する頂点の数を求める
        count_v = [0] * self.N
        for i in range(self.N):
            count_v[self.root(i)] += 1
        return count_v


n, m = map(int, input().split())
if m == 0 or n <= 2:
    exit(print(0))

UF = UnionFind(n)
ans = 0
for _ in range(m):
    u, v = map(lambda x: int(x) - 1, input().split())
    UF.unite(u, v)
    ans -= 1  # あらかじめ引いておく

for i in UF.count_root_v():
    ans += i
    ans += i * (i - 3) // 2

print(ans)