087 D - Friend Suggestions

入出力

1
2
3
4
let N,M,K = stdin.ReadLine().Split() |> Array.map int |> fun x -> x.[0],x.[1],x.[2]
let Xa = Array.init M (fun _ -> stdin.ReadLine().Split() |> Array.map int |> fun x -> x.[0],x.[1])
let Ya = Array.init K (fun _ -> stdin.ReadLine().Split() |> Array.map int |> fun x -> x.[0],x.[1])
solve N M K Xa Ya |> Array.map string |> String.concat " " |> stdout.WriteLine

解説

練習も兼ねて(破壊的な)Union-Findを簡易実装して対応します. 提出された解答を眺めていたら, 少なくともRustではpetgraph crateがあってAtCoderでも使えるようです.

破壊的な簡易Union-Find

クラスとしての実装はUnionFind.fsxに置いてあります.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
  type UnionFind = { par: int[]; size: int[]}
  let uf = { par = Array.init N id; size = Array.create N 1 }

  let rec root x =
    if uf.par.[x] = x then x
    else let r = root uf.par.[x] in uf.par.[x] <- r; r
  let find x y = root x = root y
  let unite x y =
    let rx,ry = root x, root y
    if rx=ry then false
    else
      let large,small = if uf.size.[rx]<uf.size.[ry] then ry,rx else rx,ry
      uf.par.[small] <- large
      uf.size.[large] <- uf.size.[large]+uf.size.[small]
      uf.size.[small] <- 0
      true
  let size x = let rx = root x in uf.size.[rx]

入力の変換

直接的な友達・ブロックの隣接行列を作りながらUnion-Find木を作ります. 隣接行列は.NETのResizeArray()のAddでゴリゴリと破壊的に作ります.

1
2
3
4
5
6
7
8
  let Fa = Array.init N (fun _ -> ResizeArray<int>())
  Xa |> Array.iter (fun (a0,b0) ->
    let a,b = a0-1,b0-1
    Fa.[a].Add(b); Fa.[b].Add(a); unite a b |> ignore)
  let Ba = Array.init N (fun _ -> ResizeArray<int>())
  Ya |> Array.iter (fun (c0,d0) ->
    let c,d = c0-1,d0-1
    Ba.[c].Add(d); Ba.[d].Add(c))

Faが友達(friends), Baがブロックの配列です. 入力のXaの処理の中でunite a bをかませてUnion-Find木を作っています.

最終的な計算

公式解説通り次の量を計算します.

1
2
3
(頂点 i の連結成分のサイズ)
− (頂点 i と頂点 j が同じ連結成分に含まれて、人 i と人 j が友達関係もしくはブロック関係にある j の数)
− 1

各頂点iに対して連結成分のサイズはsize i, 友達関係にある頂点jの数はFa.[i].Countで計算できます. ブロック関係にある頂点jの数はUnion-Findのfindを使って次の処理で計算できます.

1
    let blocks = Ba.[i].ToArray() |> Array.sumBy (fun b -> if find i b then 1 else 0)

あとは各頂点ごとの計算を次のようにArray.mapで処理します.

1
2
3
4
  [|0..N-1|]
  |> Array.map (fun i ->
    let blocks = Ba.[i].ToArray() |> Array.sumBy (fun b -> if find i b then 1 else 0)
    size i - Fa.[i].Count - blocks - 1)

まとめ

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
type UnionFind = { par: int[]; size: int[]}
let solve N M K Xa Ya =
  let uf = { par = Array.init N id; size = Array.create N 1 }

  let rec root x =
    if uf.par.[x] = x then x
    else let r = root uf.par.[x] in uf.par.[x] <- r; r
  let find x y = root x = root y
  let unite x y =
    let rx,ry = root x, root y
    if rx=ry then false
    else
      let large,small = if uf.size.[rx]<uf.size.[ry] then ry,rx else rx,ry
      uf.par.[small] <- large
      uf.size.[large] <- uf.size.[large]+uf.size.[small]
      uf.size.[small] <- 0
      true
  let size x = let rx = root x in uf.size.[rx]

  let Fa = Array.init N (fun _ -> ResizeArray<int>())
  Xa |> Array.iter (fun (a0,b0) ->
    let a,b = a0-1,b0-1
    Fa.[a].Add(b); Fa.[b].Add(a); unite a b |> ignore)
  let Ba = Array.init N (fun _ -> ResizeArray<int>())
  Ya |> Array.iter (fun (c0,d0) ->
    let c,d = c0-1,d0-1
    Ba.[c].Add(d); Ba.[d].Add(c))

  [|0..N-1|]
  |> Array.map (fun i ->
    let blocks = Ba.[i].ToArray() |> Array.sumBy (fun b -> if find i b then 1 else 0)
    size i - Fa.[i].Count - blocks - 1)

let N,M,K = stdin.ReadLine().Split() |> Array.map int |> fun x -> x.[0],x.[1],x.[2]
let Xa = Array.init M (fun _ -> stdin.ReadLine().Split() |> Array.map int |> fun x -> x.[0],x.[1])
let Ya = Array.init K (fun _ -> stdin.ReadLine().Split() |> Array.map int |> fun x -> x.[0],x.[1])
solve N M K Xa Ya |> Array.map string |> String.concat " " |> stdout.WriteLine