ABC436E - Minimum Swap を Python で解いた

ABC436E - Minimum Swap
https://atcoder.jp/contests/abc436/tasks/abc436_e

「なんとなくこうだろう」ぐらいの不確かさのまま AC しました。正当性を示したいところです。

コンテスト中の動き

C問題までを13分で解き、その後D問題に挑戦。詰まって一度E問題を見て、すぐに解法を思いつけるような問題ではなさそうだったのでまたD問題に戻りました。そして運良くD問題を AC できた時点で残り 33 分。一応考えてみたところなんとか合ってそうな解法を思いつき、21 分かけて AC できました。解法は思いついたけど実装がなかなかできずに焦りました。D, E ともに運良く正解でき、1178パフォーマンスでレート +14 (1047 → 1061）という良好な結果になりました。良かったです。

考察

例1, 例3を見ながら考えてみます。

仮説① 単純なコンビネーション

5
3 1 4 2 5

例1の答えが $ 6 = {}_4{C}_2 $ なので、動かす必要のない 5 以外の 4 つの数字から 2 つを選ぶ場合の数が答えになるのかな？と想像しました。

しかし、これは例3を調べれば動かす必要のない数字が 19 個あるのに答えが $ {}_{19}{C}_2 = 171 $ になっていないので明らかな誤りです。無意味な入れ替え方法が存在します。

仮説② 入れ替えたときに一方が正しい位置に移動するなら OK

次に、入れ替えたときに一方が正しい位置に移動する場合のみ入れ替えるようにすればいいかもしれないと想像しました。例1を見ます。

5
3 1 4 2 5

• 3 と 1 の入れ替え ...... 1 が正しい位置に来る
• 3 と 4 の入れ替え ...... 3 が正しい位置に来る
• 3 と 2 の入れ替え ...... どちらも正しい位置に来ない

ところが例1の答えは 6 なので、5 が関わらない 6 通りの入れ替え操作全てが有効です。ということは 3 と 2 の入れ替えも無駄ではないということですが……。ここで、試しに 3 と 2 を入れ替えてみます。

3 1 4 2 5
↓
2 1 4 3 5

確かにこのあとは 2 と 1, 4 と 3 を入れ替えることにより計3回の操作で [1,2,3,4,5] を作ることが可能です。では、いきなり 2 と 1, 4 と 3 の入れ替えをやってみましょう。するとこうなります。

3 1 4 2 5
↓
4 2 3 1 5

結局あとでさらに 4 と 1 を入れ替えないといけなくなり、計3回の操作が必要になります。というわけで正しい位置への変更にならないような入れ替え方でも必要になることがあるとわかります。

仮説③ 数字だけでなく添字の入れ替えも考える

解法が降ってくるまでの過程

ここは解法を思いつくまでのごちゃごちゃした試行錯誤を文章化しただけの独り言なので解き方を知りたい人はこの部分に関しては読まなくてもいいです。

3 1 4 2 5
↓
2 1 4 3 5

ここで先ほど考えたこの入れ替えについて、「添字 1 と添字 4 を入れ替えている」という考え方が降ってきました。先に 2 と 3、つまり添字 1 と添字 4 を入れ替えるておくことにより、それをやらなかった場合には後でやる必要になる「4 と 1 の入れ替え = 添字 1 と添字 4 の入れ替え」が必要なくなります。ということからこのようなメモをコンテスト中に書きました。

• 数字同士を入れ替えるのは当然有効。1回の入れ替えでどちらかが正しい位置に入ればOK.
• 添字同士を入れ替えるのも有効。後で他の数字同士を入れ替えるときの助けになる。

では具体的にはどのような添字の入れ替えが有効なのでしょうか。例1を見て、一番左にある「3」に注目します。この数字 3 と、ここの添字 1 に関わる数字を全部思い浮かべてみます。

• 配列 P の中で添字が 3 になっているところ、つまり P[3] = 4
• 3 の添字は 1
• 1 が入っている位置（添字）は 2

だから 3 に関して 1, 2, 4 の数字が関わってくる……みたいなよくわからないことを考えましたが、この考え方だと1つの数字に対して関係する数字が3つしか出てこないし、そうなると例3を考えると 19 個の数字があるけどそれぞれに 3 つずつ有効な入れ替えが存在するとしたら計 57 通りにしかならないし……という風に考えました。

降ってきた解法

次に考え方を変えてみました。例3を見て、一番左にある P[1] = 15 に注目します。その次は P[15] を見ます。そしたら P[15] = 12 なのでその次は P[12] を見ます。こんな発想が突然降ってきました。それでは順番に繋いでいってみましょう。

\begin{align}
P[1] &= 15\\
P[15] &= 12\\
P[12] &= 3\\
P[3] &= 13\\
P[13] &= 8\\
P[8] &= 4\\
P[4] &= 17\\
P[17] &= 10\\
P[10] &= 16\\
P[16] &= 7\\
P[7] &= 20\\
P[20] &= 1\\
\end{align}

12個の数字を使って最初の 1 に戻ってきました。同様に残った数字についても調べてみると 5個の数字 5, 9, 14, 19, 2 が 1 つのグループになり、2個の数字 6, 11 が 1 つのグループになります。ここで、同じグループから 2 つの数字を選べばいいのでは？という思考が生まれます。$ {}_{12}{C}_2 = 66 $, $ {}_5{C}_2 = 10 $, $ {}_2{C}_2 = 1 $ これらを足し合わせれば 77 となり、例 3 の答えと合います。なんかいけそうですね。

実装

1 から N までの数を見ていき、その添字を起点として中身を次々と追いかけていきます。どの数字も1度しか見ません。見るべき先が1度見た数字になったらそこでループを終えます。たったこれだけの操作を実装するのに割と手間取ってしまいました。減っていく残り時間に焦りました。なんとか残り 12 分で提出、AC できました。

N = int(input())
P = [0] + list(map(int, input().split()))
groups = []
seen = [False for _ in range(N+1)]
tmp_group = set()
for i in range(1, N+1):
  if seen[i]: continue
  now = i
  tmp_group = {now}
  while (not seen[P[now]]):
    seen[now] = True
    next = P[now]
    tmp_group.add(next)
    now = next
  groups.append(tmp_group)
  tmp_group = set()

ans = 0
for group in groups:
  v = len(group)
  if v > 0:
    ans += v * (v-1) // 2
print(ans)

まとめ

というわけで、なんかよくわからないけど法則らしきものが見つかって正解できました。記事には僕が思いついたことをそのまま書いていったので、ヒントになりそうなことを眺めていたら突然解法が降ってきたようにしか見えなくて何の参考にもならない記事になってしまったと思います。実際本当に降ってきただけなので何も書きようがないのですが……。

正当性の確認

さて、この解法の正当性を考えてみます。なぜ合ってるのか。ChatGPTに相談していくうちにわかったようなわからないような理屈が得られたのでそれを書いていきます。まず、先ほど書いた

\begin{align}
P[1] &= 15\\
P[15] &= 12\\
P[12] &= 3\\
......
\end{align}

という操作がありますが、公式解説でいわれるようにこれをグラフに見立てます。1 から N までの頂点があり、 i → P[i] という有向辺がそれぞれあるわけですね。この頂点 i から伸びる辺の向かう先をしかるべき頂点、具体的には i 自身につけかえるのがこの問題の目的です。全ての頂点をこのようにします。つまり、いくつかある閉路を全て分解して独立させるのが目的です。

この問題の P[i], P[j] を入れ替える操作はグラフの中から辺を 2 本選んでそれぞれの向かう先を交換することに相当します。そして、同じ閉路に所属する辺 2 本を入れ替えることによって 2 つの閉路に分解されます。どのように選んでも分解されますし選び方によっては長さ 1 の閉路すなわち目的を達成した頂点が生まれます。

このようにして閉路を小さく分けていき、最終的に全ての閉路を分解して全てを独立した頂点にしてしまうのが目的です。ですから同じ閉路の中のどの 2 本の辺を選んでも有効な操作となります。

一方、違う閉路同士の辺を選んでしまうと 2 つの閉路がくっついて大きな 1 つの閉路になってしまい、目的から遠ざかります。

このような理由から、同じ閉路に所属する 2 辺を選べば目的を達成できる操作となり、異なった閉路に所属する 2 辺を選べば目的を達成できない操作となります。というわけで初手の数 = 同じ閉路に所属する 2 辺を選ぶ場合の数です。

おまけ

ちなみにコンテスト中の提出コードはもっとぐちゃぐちゃです。開始点の候補をキューにして入れておいて……みたいなことを考えていたせいで無駄に複雑になっています。グループ数の計算もなぜか1小さい値になってしまうので最後の計算式で v を v+1 に置き換えて帳尻合わせをしています。とりあえず AC できればいいやというその場しのぎですがうまくいって助かりました。

from collections import deque
N = int(input())
P = [0] + list(map(int, input().split()))
groups = []
seen = [False for _ in range(N+1)]
tmp_group = set()
for i in range(1, N+1):
  if seen[i]: continue
  que = deque([i])
  while(que):
    q = que.pop()
    seen[q] = True
    if seen[P[q]]:
      groups.append(tmp_group)
      tmp_group = set()
    else:
      tmp_group.add(q)
      que.append(P[q])

ans = 0
for group in groups:
  v = len(group)
  if v > 0:
    ans += v * (v+1) // 2
print(ans)