090 D - Even Relation

入出力

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

解説

DFSで素直に木を走査します.

隣接リスト生成

二点間の距離の処理によって全体の処理が変わるため, そこに注意すればあとは素直に作るだけです. ここではw&&&1とビットの論理和で処理します.

1
2
3
4
  let Aa =
    (Array.init N (fun _ -> []),Ia)
    ||> Array.fold (fun Aa (u,v,w) ->
      Aa.[u-1]<-(v-1,w&&&1)::Aa.[u-1]; Aa.[v-1]<-(u-1,w&&&1)::Aa.[v-1]; Aa)

他の問題ではResizeArrayで処理したときもありますが, ここではListで処理しています. 今回せっかくなので試してみたら, 少なくとも今回のケースではListの方が速いようでした.

DFS

再帰とfoldで隣接リストを走査します. 今回はArray.zeroCreate Nで初期化した配列の値を再帰の中でゴリゴリ書き換える形で実装します.

根と初期値を適当に決める必要があります. ここではdfs関数をpiを根のインデックス, ciを子ノードのインデックス, vを0,1の値として構成します. 特に次の形で計算をはじめます.

1
2
  let rec dfs pi ci v Xa = "処理を書く"
  Array.zeroCreate N |> dfs 0 0 0

dfsを実装しましょう. 本体のfold処理は次のように書きます.

1
2
3
4
5
  let rec dfs pi ci v Xa =
    // 値の書き換え処理が必要
    Array.get Aa ci
    |> List.filter (fun (i,_) -> i <> pi)
    |> List.fold (fun acc (gci,w) -> dfs ci gci (v^^^w) acc) Xa

まずArray.get Aa ciでdfsの引数で指定した隣接リストの子ノードを取ります. Array.filterで親のノードを除外した上でfoldの中でXaをゴリゴリ書き換えます. foldのgciはdfsに「子ノードの子ノード」として食わせる前提で孫(grandchild)の想定で名付けた変数です. 再帰的にdfs ci gciで呼び出して値を更新します. ここでv^^^wと排他的論理和で距離を処理しています. 親子で偶奇が違う場合だけフラグが立ちます.

最後にコメントアウトで残した書き換え処理を考えましょう. ciで渡した子ノードの値を書き換えます. ここではv^^^1として排他的論理和で値を書き換えます. 結果的にdfsは次のように書けます.

1
2
3
4
5
6
  let rec dfs pi ci v Xa =
    Array.set Xa ci (v^^^1)
    Array.get Aa ci
    |> List.filter (fun (i,_) -> i <> pi)
    |> List.fold (fun acc (gci,w) -> dfs ci gci (v^^^w) acc) Xa
  Array.zeroCreate N |> dfs 0 0 0

まとめ

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
let solve N Ia =
  let Aa =
    (Array.init N (fun _ -> []),Ia)
    ||> Array.fold (fun Aa (u,v,w) ->
      Aa.[u-1]<-(v-1,w&&&1)::Aa.[u-1]; Aa.[v-1]<-(u-1,w&&&1)::Aa.[v-1]; Aa)

  let rec dfs pi ci v Xa =
    Array.set Xa ci (v^^^1)
    Array.get Aa ci
    |> List.filter (fun (i,_) -> i <> pi)
    |> List.fold (fun acc (gci,w) -> dfs ci gci (v^^^w) acc) Xa
  Array.zeroCreate N |> dfs 0 0 1

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