OCaml で AtCoder のコンテストに出始めて約1年が経ちました。 知見が溜まってきたのでまとめておきます。

注意

本記事の想定環境は2023年8月の言語アップデート後の AtCoder の環境です (参考。2024年5月18日に確認)。 例えば、OCaml のバージョンは 5.0.0 を想定しています。

Scanf.scanf では %d や %s の前に半角スペースを付ける

C言語のscanfは%dや%sだけでも値を読み込む前に空白文字 (スペース、タブ、改行)を読み飛ばしますが1、 OCaml のscanfはそれだけだと読み飛ばしません。読み飛ばすことを明示する必要があります。

%dや%sの前に (半角スペース)を付けるとすべての種類の空白文字を読み飛ばします。また、\nを付けると改行だけを読み飛ばせます。 ただ、自分の経験では、AtCoder で値の区切り文字がどの空白文字なのかを気にしたことはないので、 だけで十分だと思います。

加えて、とりあえず何にでも を付けておくといいと思います。 そうすることで、次のようなミスを防げます。

(* ケース1: "1 2"のように1行に2つの値が書かれていて、それを1つずつ読む *)
(* x を読んだ後、スペースを読み飛ばさないのでエラーになる *)
let () =
  let x = Scanf.scanf "%d" Fun.id in
  let y = Scanf.scanf "%d" Fun.id in
  Printf.printf "%d, %d\n" x y

(* ケース2: 1行に1つ値が書かれていて、それを1つずつ読む *)
(* 1行目を読んだ後、改行を読み飛ばさないのでエラーになる *)
let () =
  for i = 1 to 5 do
    let x = Scanf.scanf "%d" Fun.id in
    Printf.printf "%d\n" x
  done

ただし、例外もあります。例えば、空文字列を読む必要があるときに %sを使ってしまうと意図しない文字列が読み込まれてしまいます。

Printf.printf で flush する

%!でできます。インタラクティブな問題で必要になります。

64ビット符号付き整数が必要な問題が存在する

OCaml のintは最大でも63ビット符号付き整数の型なので、答えが64ビット符号付き整数にギリギリ収まる問題ではintを使うと WA になります。

このような問題では代わりにint64を使います。 int64の基本的な演算はOCaml library : Int64に書かれています。 また、int64の値をScanf.scanfで読み込むときは%Ldを指定します。

Set と Map の cardinal は Θ(1) ではない

Set.Make、Map.Makeによって生成されるモジュールにあるcardinal関数は Θ(要素数) です。このことは実装を見ると分かります (ドキュメントには明記されていません)。

そのため、要素数が何回も必要になる問題で使うと TLE になります。

もし、要素数を Θ(1) で取得したいときは core のCore.Mapモジュールを使うと良さそうです (自分はまだ使ったことがありません)。Core.Map.lengthがcardinalに相当します。

length map returns the number of elements in map. O(1), but Tree.length is O(n). (引用元: core v0.16.1 · OCaml Package)

また、もし、要素間の大小関係に基づく操作が必要ないときは標準ライブラリのHashtblモジュールも使えます。Hashtbl.lengthがcardinalの代わりになります。

補足: cardinal関数を Θ(1) にしたいという要望はすでに GitHub の issue にあります。気になる方はMake Map length run in O(1) instead of O(n) · Issue #12866 · ocaml/ocamlをご参照ください。

iter を多用しすぎると遅くなることがある

(余力ができたら追記します)

うろ覚えなのですが、以前 iter を多用したプログラムを書いたらかなり遅く、 それを for 文で書き直したら早くなったことがありました。

おそらく以下の提出のことだと思います (時系列が逆になっていますが、これは最初に書いた iter を使ったプログラムがうまく動かず、後で書き直したからです)。

簡単にしか調べていない、かつ、調べたのがかなり前なので不確かなのですが、 AtCoder の環境では flambda が有効になっていないために、iter を使ったプログラムの最適化が弱いようです。 上のツイートによると、flambda を有効にしたら2、時間差が縮まったそうです。

TLE が厳しい問題では iter の多用しすぎにお気を付けください。

処理の途中で出てくる値をイテレーターの形で集める

例えば、グラフを探索しているときに訪れた頂点を記録したいことがあります。 そういうときは次のように可変変数に格納するのが一般的だと思います。

let visited_nodes = ref []
let rec dfs v =
  visited.(v) <- true;
  visited_nodes := v :: !visited_nodes;
  List.iter (fun w ->
    if not visited.(w) then dfs w
  ) al.(v)

let () =
  dfs 1;
  (* 訪れた順にするには反転する必要がある *)
  visited_nodes := List.rev !visited_nodes;
  ...

しかし、次のようにイテレーターの形で頂点を集めることもできます。

(* 集めた頂点を後で処理するための関数 k を引数に加える *)
let rec dfs v k =
  visited.(v) <- true;
  (* v を訪れたので k に v を渡す *)
  k v;
  List.iter (fun w ->
    if not visited.(w) then dfs w k
  ) al.(v)
(* iter で使えるようにする *)
let dfs v = Iter.from_iter (dfs v)

let () =
  let visited_nodes = dfs 1 in
  ...

このように書くと次のような利点があります。

  • 値を集める処理と値を使う処理を分けて書ける。そうすると、二つの処理を同時に考えずに済むので、個人的には頭に優しいと思います。
  • 値が得られた順に並ぶ。上の可変変数の例では最後にリストを反転させる必要がありました。
  • メモリにすべての値を置かないで処理が行なえ、速度の向上が期待できる。リストや配列の形で保持、操作すると最初の値はもちろん、中間結果もメモリに格納する必要があります。

この方法を使ったプログラムの例をいくつか載せておきます。

毎回決まった量加算するループを iter を使って書く

C++ で言うところのfor(int i = l; i < r; i += step){ ... }は iter を使って次のように書けます3。

Iter.iterate ((+) step) l
|> Iter.take_while (fun i -> i < r)
|> Iter.iter (fun i -> ...)

iter にはIter.int_range_byというぴったりそうな関数がありますが、上のプログラムと異なる動作をするケースがあるので注意が必要です。

let step = 2 in
let l = 0 in
let r = -1 in
(* for(int i = l; i < r; i += step) のつもり *)
(* 何も出力しないでほしいが、実際は 0 を出力する *)
Iter.int_range_by ~step l r
|> Iter.iter (Printf.printf "%d\n")

(step > 0で)l >= rのときは何もしないように条件分岐すれば問題ありませんが、焦っていると忘れてしまうかもしれないので、個人的には最初の方法で書くのがオススメです。

未定を None で表現するときの最小値、最大値

最小値 (または最大値)を動的計画法で求めたいとします。 テーブルの初期値にはmax_intや-1などの具体的な値を指定することが多いですが、そうすると途中の計算がおかしくならないか気を配る必要があります。 そのため、代わりにNoneを指定して未定とそれ以外を明確に区別することにします。 このとき、テーブルの値は'a optionになるわけですが、'a optionの最小値 (または最大値)はどのように計算すればいいでしょうか。

最大値

(注意: 以下では実行速度のことはあまり考えていません。記述量だけを考慮しています)

少なくとも AtCoder の環境ではNone <= Some vで、Some v <= Some wとv <= wが同値なので、単にmaxと書けます4。

最小値

単にminと書くことはできません。色々試してきましたが、結局のところ、次のような関数をあらかじめ定義しておくのが良さそうです5。

let lift ~f o1 o2 =
  match o1, o2 with
  | None, _ -> o2
  | _, None -> o1
  | Some x1, Some x2 -> Some (f x1 x2)

このliftを使うと、2つの値の最小値を取る関数はlift ~f:minと書けます。

この方法は最大値でも使えます。両方で同じ方法を使えば書き方を迷う時間を減らせます。 さらには、一般の二項演算でも使えて汎用性が高いです。

特定の型に特化したハッシュテーブルを使う

多相的なハッシュテーブル (Hashtbl.t)だと遅くて困るときは使用する型に特化したハッシュテーブルを使うと速くなるかもしれません。

例えば、キーがintのハッシュテーブルは

Base.Hashtbl.create ~size:n (module Base.Int)

で作成できます。

また、Hashtbl.Makeを使っても作成できます。

(* バージョン 5.1 以降では Int.hash があるので単に Int と書けます。 *)
module Int_tbl = Hashtbl.Make (Base.Int)

Int_tbl.create n

AtCoder Beginner Contest 363 C - Avoid K Palindrome 2の解法で簡単に試したところ、

  • Base.Hashtblを使ったプログラムが一番速かった。
  • Hashtbl.Makeを使ったプログラムはその約1.65倍。
  • 普通のHashtblを使ったプログラムは約2.28倍。

(もちろん精査は必要ですが、)速度に困ったらとりあえず初めにBase.Hashtblを使えば良さそうです。

参考資料

  1. これが仕様としてどこかに明記されているかは不明です。 ↩︎

  2. 記憶違いでなければ、ocaml-base-compilerの代わりにocaml-variants.4.14.1+optionsとocaml-option-flambda.1をインストールしたはずです。これらのパッケージについては参考資料をご覧ください。 ↩︎

  3. 元のコードも OCaml のコードもl、r、stepがオーバーフローが起こらない程度に小さいと暗に仮定していることにお気を付けください。例えば、OCaml だとl = 0, r = max_int, step = max_int - 1のとき、オーバーフローが起こり、iが負の数になります。 ↩︎

  4. これが仕様としてドキュメントに明記されているかは不明です。 ↩︎

  5. (fが結合律を満たすとき、)lift ~fはfを演算に持つ半群に単位元を追加して得られるモノイドの演算です。 Haskell を使っている方にはお馴染みかもしれません。 ↩︎