093 D - Transit Tree Path

解説1, DFS

公式解説通りにDFSで実装します.

入出力

1
2
3
4
5
let N = stdin.ReadLine() |> int
let Aa = Array.init (N-1) (fun _ -> stdin.ReadLine().Split() |> fun x -> int x.[0], int x.[1], int64 x.[2])
let Q,K = stdin.ReadLine().Split() |> Array.map int |> (fun x -> x.[0],x.[1])
let Xa = Array.init Q (fun _ -> stdin.ReadLine().Split() |> Array.map int |> fun x -> x.[0],x.[1])
solve N Aa Q K Xa |> Array.iter stdout.WriteLine

距離cだけint64にしています.

隣接リストの生成

ここまでに何度か出てきたのと同じように素直に実装すれば問題ありません.

1
2
3
  let Ga =
    (Array.init N (fun _ -> []), Aa)
    ||> Array.fold (fun Ga (a,b,c) -> Ga.[a-1] <- (b-1,c)::Ga.[a-1]; Ga.[b-1] <- (a-1,c)::Ga.[b-1]; Ga)

DFSの実装

Array.zeroCreate N |> dfs (K-1) (-1) 0Lを前提に実装すると次のようにすっきり書けます.

1
2
3
  let rec dfs v p d Da =
    Array.set Da v d
    (Da, Ga.[v]) ||> List.fold (fun Da (u,c) -> if u=p then Da else dfs u v (d+c) Da)

F#の配列は破壊的なデータ型で, Array.setはunitを返すだけで書き換えた配列を返してくれるわけではないため, Array.set Da v dを別立てにしています. foldの中身は隣接リスト内の値が親ノードと一致していたら積んだ値を返し, それ以外は入力のdに対して隣接リストが持つ距離cを素直に足して積むだけです.

二点間の距離の計算

あとはKからの最短距離を素直に計算します.

1
  Xa |> Array.map (fun (x,y) -> Da.[x-1]+Da.[y-1])

まとめ

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
let solve N (Aa:(int*int*int64)[]) Q K Xa =
  let Ga =
    (Array.init N (fun _ -> []), Aa)
    ||> Array.fold (fun Ga (a,b,c) -> Ga.[a-1] <- (b-1,c)::Ga.[a-1]; Ga.[b-1] <- (a-1,c)::Ga.[b-1]; Ga)
  let rec dfs v p d Da =
    Array.set Da v d
    (Da, Ga.[v]) ||> List.fold (fun Da (u,c) -> if u=p then Da else dfs u v (d+c) Da)
  let Da = Array.zeroCreate N |> dfs (K-1) (-1) 0L
  Xa |> Array.map (fun (x,y) -> Da.[x-1]+Da.[y-1])

let N = stdin.ReadLine() |> int
let Aa = Array.init (N-1) (fun _ -> stdin.ReadLine().Split() |> fun x -> int x.[0], int x.[1], int64 x.[2])
let Q,K = stdin.ReadLine().Split() |> Array.map int |> (fun x -> x.[0],x.[1])
let Xa = Array.init Q (fun _ -> stdin.ReadLine().Split() |> Array.map int |> fun x -> x.[0],x.[1])
solve N Aa Q K Xa |> Array.iter stdout.WriteLine

解説2: ダイクストラ法

他のF#解答を参考にダイクストラ法での実装も紹介します. もとのコードは二つポイントがあります.

もとの入力をリストで処理している一方, 私は配列で実装しています. 処理系も違うため単純な速度比較はできません. どなたか入力を合わせて私のコードと速度比較してみてください. そのうち自分で実装してみようとは思っています.

入出力

1
2
3
4
5
let N = stdin.ReadLine() |> int
let Aa = Array.init (N-1) (fun _ -> stdin.ReadLine().Split() |> fun x -> int x.[0], int x.[1], int64 x.[2])
let Q,K = stdin.ReadLine().Split() |> Array.map int |> (fun x -> x.[0],x.[1])
let Xa = Array.init Q (fun _ -> stdin.ReadLine().Split() |> Array.map int |> fun x -> x.[0],x.[1])
solve N Aa Q K Xa |> Array.iter stdout.WriteLine

大枠

DFSと同じで次のように書けます.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
let solve N (Aa:(int*int*int64)[]) Q K Xa =
  let Ga =
    (Array.init N (fun _ -> []), Aa)
    ||> Array.fold (fun Ga (a,b,c) -> Ga.[a-1] <- (b-1,c)::Ga.[a-1]; Ga.[b-1] <- (a-1,c)::Ga.[b-1]; Ga)
  let Da = dijkstra N (K-1) Ga
  Xa |> Array.map (fun (x,y) -> Da.[x-1]+Da.[y-1])

let N = stdin.ReadLine() |> int
let Aa = Array.init (N-1) (fun _ -> stdin.ReadLine().Split() |> fun x -> int x.[0], int x.[1], int64 x.[2])
let Q,K = stdin.ReadLine().Split() |> Array.map int |> (fun x -> x.[0],x.[1])
let Xa = Array.init Q (fun _ -> stdin.ReadLine().Split() |> Array.map int |> fun x -> x.[0],x.[1])
solve N Aa Q K Xa |> Array.iter stdout.WriteLine

つまりDFSがdijkstraに変わっただけです.

ダイクストラ法

私はまだ一般的なダイクストラ法をきちんと理解できていません. 詳しくはアルゴリズムの本を参照してください.

今回の実装に関しては次の通りです.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
let dijkstra N K (Ga: (int * int64) list []) =
  let Da = Array.create N System.Int64.MaxValue
  Da.[K] <- 0L
  let rec loop (Da:int64[]) q =
    if Set.isEmpty q then (Da,q)
    else
      let (c,v) = Set.minElement q
      let q0 = Set.remove (c,v) q
      if c <= Da.[v] then
        ((Da,q0), Ga.[v]) ||> List.fold (fun (Da,q) (bi,ci) ->
          let s = Da.[v]+ci
          if s < Da.[bi] then Da.[bi] <- s; (Da, Set.add (s,bi) q) else (Da,q))
      else (Da,q0)
      |> fun (Da,q) -> loop Da q
  loop Da (Set.singleton (0L,K)) |> fst

もとのコードはloopが完全に破壊的です. このコードもDa.[K] <- 0Lやloopの中でのDa.[bi] <- sが厳密には破壊的なコードです. わざわざ非破壊的にするほどでもないため, 大目に見て実装しています.

ポイントは優先度つきキューの代わりにSetを使っている点です. キュー代わりのSetに積んだ値から最小値を取り出し, キューが尽きるまでループをくり返しています. もちろんSetでは速度は出ません. また配列のArray.setと違ってSet.addは更新したSetを返してくれる非破壊的な関数です.

ちなみに.NET6で優先度つきキューが実装されたものの, AtCoder上のF#は.NET Core 3.1.201で使えません.

TODO 確認: 解説2とオリジナルの破壊的な実装の速度比較