070 C - Snuke Festival

自分用の記録

公式解説では二分探索を提案しています. それぞれ次のような解答例があります.

F#版はstart-1やstop+1を外して書けないかと思って轟沈し, Haskell版も番兵なしで書けないかと思って轟沈し, 自力実装もTLEではまり倒したため, CやPythonなどのコードを見つついろいろ試し, 最終的にはRustのコードを参考にしました. ここでもこれに基づいた解説をつけます.

解説

参考にした命令型のRustのコードの焼き直しを説明したあと, 関数型に書き換えたバージョンを説明します.

まずは公式解説と同じく配列で入力を取ってソートします. あとの都合があるためBaもソートします.

1
2
3
  let Xa = Aa |> Array.sort
  let Ya = Ba |> Array.sort
  let Za = Ca |> Array.sort

以下ソートした配列Xa, Ya, Zaで考えます. 何も考えずRustを直移植した結果は次の通りです.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
  let Xa = Aa |> Array.sort
  let Ya = Ba |> Array.sort
  let Za = Ca |> Array.sort
  let mutable i = 0
  let mutable k = 0
  let mutable ans = 0L
  for b in Ya do
    while i<N-1 && Xa.[i+1]<b do i <- i+1
    while k<N-1 && Za.[k]<=b  do k <- k+1
    if Xa.[i]<b && b<Za.[k] then ans <- ans + (int64 (i+1) * int64 (N-k))
  ans

アルゴリズム上のポイントはBaのソートとi, kの引き継ぎです. AaとCaをソートしているため配列の引数からそのまま個数が計算できます. さらにBaのソートのおかげで, Baに関する添字をjとすれば, j+1に関する結果を計算するときjで計算したiとkから再開でき, 余計な計算が省略できます. Zaの関する計算は問題の素直な不等式ではなくZa.[k]<=bにした上でN-kを計算しています. (解説を書きながら気づいたのですが, Array.sortDescendingを使えば問題の条件のまま計算できます.)

F#コードのポイントはans <- ans + (int64 (i+1) * int64 (N-k))にもあります. ここでint64 (i+1) * (N-k)としてしまうとWAが出ます. 実際に出てはまり倒しました. 理由はintで(i+1) * (N-k)が計算されるとオーバーフローを起こす場合があり, このかけ算自体もint64で計算しないといけないからです. 配列の添字はintでint64にはできないため, 配列の添字を使った計算には注意が必要です.

命令型での焼き直しができました. 一般的には命令型の方が速くメモリ効率もよいため, これで終わらせても問題ありません. しかし関数型競技プログラミングを標榜している以上やはり関数型に書き換えます.

まずはfor b in Ya doから書き換えます. 最終的にはlet mutable ans相当の値を返したいため, foldを使って処理すればいいでしょう. 具体的には次のように書けます..

1
2
3
4
5
6
7
  let mutable i = 0
  let mutable k = 0
  let mutable ans = 0L
  (0L, Ya) |> Array.fold (fun acc b ->
    while i<N-1 && Xa.[i+1]<b do i <- i+1
    while k<N-1 && Za.[k]<=b  do k <- k+1
    if Xa.[i]<b && b<Za.[k] then acc + (int64 (i+1) * int64 (N-k)) else acc)

F#では最後に計算した式の値を返してくれるため, これを関数の最後に書けば問題ありません. ポイントはifに対してelse節を書いて値を返す部分です. そもそもelse節がないとエラーを吐いてくれるため, そこで気付けるでしょう.

あとはmutableのiとkをうまく処理する必要があります. 積んだ結果を引き回す必要があるため, accをタプルに変えて積む必要があります.

ループは一般に再帰で書き換えられます. これもelse節で値を返す必要があり, 特に次のように書けます.

1
2
  let rec a i b = if i<N-1 && Xa.[i+1]<b then a (i+1) b else i
  let rec c k b = if k<N-1 && Za.[k]<=b  then c (k+1) b else k

関数名は一文字でaとcです. 短いプログラムだからこれで十分意図は伝わるでしょう. 気にいらなければもう少し長く説明的な名前でも構いません. この関数を使うとfoldは次のように書き換えられます.

1
2
3
4
5
  ((0L,0,0),Ya) ||> Array.fold (fun (acc0,i0,k0) b ->
    let i = a i0 b
    let k = c k0 b
    let acc = if Xa.[i]<b && b<Za.[k] then acc0 + (int64 (i+1) * int64(N-k)) else acc0
    (acc,i,k)) |> fun (acc,_,_) -> acc

accが(acc,i,k)に書き換わりました. 計算結果もタプルで返します. 先程と違い, foldの最終的な返り値もタプルになってしまうため, 最後に必要な第一要素だけ取り出す関数を噛ませてあります. 最終的な全体実装は次の通りです.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
let solve N Aa Ba Ca =
  let Xa = Aa |> Array.sort
  let Ya = Ba |> Array.sort
  let Za = Ca |> Array.sort
  let rec a i b = if i<N-1 && Xa.[i+1]<b then a (i+1) b else i
  let rec c k b = if k<N-1 && Za.[k]<=b  then c (k+1) b else k
  ((0L,0,0),Ya) ||> Array.fold (fun (acc0,i0,k0) b ->
    let i = a i0 b
    let k = c k0 b
    let acc = if Xa.[i]<b && b<Za.[k] then acc0 + (int64 (i+1) * int64(N-k)) else acc0
    (acc,i,k)) |> fun (acc,_,_) -> acc

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

コードをもっと短くしたければ次のように読み込み時点でソートできます.

1
2
3
let Aa = stdin.ReadLine().Split() |> Array.map int64 |> Array.sort
let Ba = stdin.ReadLine().Split() |> Array.map int64 |> Array.sort
let Ca = stdin.ReadLine().Split() |> Array.map int64 |> Array.sort