OCamlでTypical Dynamic Programming Contest E問題を解いてみた

ここ一ヶ月ほどAtCoderをPythonですすめてきて、Pythonの遅さにひやりとすることはあっても基本的にアルゴリズムがあっていて素直な最適化を施していればABCの問題であれば通せそう、という感触を得ている。

しかし少し上の問題を見てみるとそうはいかないようだ。例えばAtCoderのTypical Dynamic Programming ContestというDP問題ばかり集めたサンプルコンテストのE問題：

tdpc.contest.atcoder.jp

かなり素直な桁DPなのだが、桁数が10000でmodをとる値が最大100なので計算量のオーダーが10000 * 100 * 10（最後の10は0~9までの数で回すから）、つまり10⁷になる。 この計算量はコンパイル言語なら余裕だがPythonだと普通にアウトのようだ。

tdpc.contest.atcoder.jp

（まあこのコードが定数で最適かといわれると違うのだが・・・）

というわけでコンパイル言語で書いてみる。

競技プログラミング的に素直な選択はC++だが、もう少しPythonに近い宣言的な表現力のある言語がいい。

表現力という意味ではHaskellが相当高いが、純粋関数型だとアルゴリズムにけっこう工夫が必要になり、Pythonで考えるのとはまったく違う世界になる。

速度もあり、適度に関数型で宣言的に記述することができ、かつ必要なところで副作用を気軽に扱える言語、ということでOCamlを選んだ。

tdpc.contest.atcoder.jp

最長251msで無事AC。

open Batteries
 
let digitList s = s |> String.to_list |> List.map (fun c -> Char.code c - 48)
let id x = x
 
let d = Scanf.scanf "%d\n" id
let n = Scanf.scanf "%s\n" id
let m = 1000000007
 
let f x n =
  let v, a = x in
  let a1 = Array.make d 0 in
  for i = 0 to d-1 do
    for j = 0 to 9 do
      a1.((i+j) mod d) <- (a1.((i+j) mod d) + a.(i)) mod m
    done
  done;
  for j = 0 to n-1 do
    a1.((v+j) mod d) <- (a1.((v+j) mod d) + 1) mod m
  done;
  ((v+n) mod d, a1)
 
let v, a = List.fold_left f (0, Array.make d 0) (digitList n)
let () = Printf.printf "%d\n" (a.(0)  + (if v == 0 then 1 else 0) - 1)

まともにOCamlを書くのははじめてだったのでまだいろいろ汚いと思うが、とりあえず自分がしたいことを比較的素直に記述できて速度が出るのは非常にうれしい。

f関数の中で動的計画法らしくArrayに対する破壊的代入を繰り返しているが、引数には変更は加えずに命令的に作成したArrayを結果として返しているので外側からみると純粋関数だ。

そのf関数を引数としてfold_leftで畳み込んだり、文字列を|>を使った関数のチェーンで一桁の数字のリストに変換したり、といったところは抽象度が高い関数型らしい記述ができる。

その上で実行速度が10⁷のオーダーをAtCoderの時間制限の2000msでも問題なくこなせるのだからこれは強い。

習熟度が非常に低いので実際のコンテストではまだ使えないだろうが、少しずつ過去問（とくに計算量のオーダーが大きそうな問題で）をOCamlで挑戦していきたい。