089 E - Colorful Hats 2

入出力

1
2
3
let N = stdin.ReadLine() |> int
let Aa = stdin.ReadLine().Split() |> Array.map int64
solve N Aa |> stdout.WriteLine

解説

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

MOD計算

計算漏れしないように演算子を定義してそれを使いましょう.

1
2
  let MOD = 1_000_000_007L
  let (.*) a b = (a*b)%MOD

MODはInt32でも問題ないと思いますが, たまにオーバーフローしてはまり倒すため, 怪しそうな場合はとにかくInt64に倒します.

大枠

各iごとに計算した結果をかけて積んでいけばよく, 単純にArray.foldでループを回します. 変数名は公式解説に合わせます.

積む変数は最終的な計算用の値であるTiとxi,yi,ziです. 前の人までの帽子の値が必要なためxi,yi,ziもStateに積む必要があります. 初期値はかけ算の初期値だからTi = 1, 帽子の数は(xi,yi,zi) = (0,0,0)でよく, 特に次のように書けます.

1
2
3
  ((1L,(0L,0L,0L)), Aa)
  ||> Array.fold (fun ((t,(x,y,z)) ai) -> ("適切に埋める")
  |> fst

最終的に必要なのはtの値だからそれをfstで切り取ります. タプルを切り取るのはfstとsndまでしかないため, この最後の切り取りが単純になるようにStateを構成しています.

folderの構成

まずtを更新します. filterとlengthを組み合わせるのが素直な実装です. ここではsumByで次のように処理します.

1
    let ti = [|x;y;z|] |> Array.sumBy (fun w -> if w=ai then 1L else 0L)

いまはInt64でtを積むため, filter >> lengthで処理する場合は最後にint64をかませる必要があります.

次は(x,y,z)の更新です. これは単純な場合分けで十分です.

1
    let xyz = if ai=x then (x+1L,y,z) elif ai=y then (x,y+1L,z) else (x,y,z+1L)

あとはこれらの値をタプルにまとめて次のステップに回します.

1
    (t.*ti, xyz)

MODつきのかけ算にするよう注意しましょう.

まとめ

ここまでの処理をまとめると次のように書けます.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
let solve N Aa =
  let MOD = 1_000_000_007L
  let (.*) a b = (a*b)%MOD
  ((1L,(0L,0L,0L)), Aa)
  ||> Array.fold (fun (t,(x,y,z)) ai ->
    let ti = [|x;y;z|] |> Array.sumBy (fun w -> if w=ai then 1L else 0L)
    let xyz = if ai=x then (x+1L,y,z) elif ai=y then (x,y+1L,z) else (x,y,z+1L)
    (t.*ti, xyz))
  |> fst

let N = stdin.ReadLine() |> int
let Aa = stdin.ReadLine().Split() |> Array.map int64
solve N Aa |> stdout.WriteLine