塩見周子の徒然日記

自分のことを塩見周子と思い込んでいるオタクです

3/4 AtCoder(サイズ付きUnion Find)

こんばんは。昼過ぎに起きて人生が終了しがちです。とーです。

AtCoder
解くに至った問題(1問)
ABC:120D

ABC120D
島の間に橋が架かっていて、これを一本ずつ落としていったら、いくつかの橋を使って互いに行き来できなくなる島のペアの数を列挙していこうねっていう趣旨の問題

昨日の記事に思考の流れを書いた通り、UF木を扱ううえで「枝を切る」という操作はやりにくい(というか実装法を知らない)ので、逆に後ろの方から橋を架けていくというやり方を採用していく

さて、橋を架けるという操作を扱う時、この問題で考えなければならないのは「橋を架けることで新たに行き来可能になる島は何通りあるか」ということである

結論から言ってしまえば、
①橋を架けた島同士が、UF木で同じ根を持っていた場合(つまり同じグループ)、新たに行き来できるようになる島はない
②橋を架けた島同士が、UF木で違う根を持っていた場合、新たに行き来できるようになる島のペア数は、その根を持つ島(繋いだ島と同じグループに属する島)の数を掛け合わせたものになる(例として、根1に島が二つ(ア、イ)、根2に島が3つ(ウ、エ、オ)属していたら、イとウをつなぐことで、この二つの島を介してア、イ、ウ、エ、オは全て行き来できるようになり、新たに行き来できるようになった島のペア数は2*3 == 6)

この二つだけを念頭に置けばよい

問題を解くうえで、「根に島がいくつ属しているか」という情報が必要になる
今回はsizeという値で管理することにする sizeは、最初どの島も1を持っているが、UF木で言うところのuniteの操作をするたびに、吸収した側の島の値に、吸収された側の根のsizeを加えていくことで、共通の根を持つグループの要素数を管理することができる

以下にコードを書いていく
まずは入力N,Mと橋の情報を格納するリストL、そしてサイズ付きUFの実装を行う

N,M = map(int,input().split())
L = []
for i in range(M):
  L.append(list(map(int,input().split())))
L.reverse()         #橋が0の状態から架けていくので、リストを逆順にする(この順番で橋を架けていく)
par = []
rank = [0]*N
size = [1]*N           #サイズの情報を加えたUFを作っていく
for i in range(N):
  par.append(i)
def find(x,par):
  if x == par[x]:      
    return x
  else:
    return find(par[x],par)
def unite(x,y,par,rank,size):       #ここもsizeを加えて改良する
  x = find(x,par)
  y = find(y,par)          
  if x != y:
    if rank[x] < rank[y]: 
      par[x] = y
      size[y] += size[x]      #sizeの変更方法は上で説明した通り
    else:
      par[y] = x
      size[x] += size[y]
      if rank[x] == rank[y]:
        rank[x] += 1
def same(x,y,par):
  return find(x,par) == find(y,par)

ここから、「入力で与えられた島をつなげると、新たに行き来できる島の組がいくつできるか」を調べていく 数え方は上で述べた

res = []      #ここに「新しく行き来できるようになった島のペア数」を保管する
for i in range(M):
  A = find(L[i][0]-1,par)     #入力は0インデックスに注意
  B = find(L[i][1]-1,par)
  if A == B:
    res.append(0)        #同じ根を持っていたら新たに行き来できる島のペアは増えない resに0を入れる
  else:
    res.append(size[A]*size[B])     #違う根であれば、sizeで保管した要素数を掛け合わせる
    unite(A,B,par,rank,size)      #ペア数をresに入れたら、島同士を結ぶ
ans = 0
for i in range(M):
  ans += res[M-i-1]      #行き来できるようになった島を数えたので、逆順に出力すれば「橋を落とすことで行き来できなくなる島の数」になる
  print(ans)


atcoder.jp

今回もじゅっぴーダイアリーを参考にさせていただきました はてな記法の件もこの場を借りてお礼を申し上げます ありがとうございました
juppy.hatenablog.com