読者です 読者をやめる 読者になる 読者になる

チラシの裏

プログラミングとか色々

OCamlで車輪を再発明する

最近研究でCoqぐらいしか書いてないので、リハビリがてらAtCoderで 開かれるコンテストにOCamlで参加して遊んでます。 僕は個人的な趣味もあってOCamlを使う際は専ら言語処理系を書いていましたが、 競技プログラミングでは全く関係ない題材を扱う必要があるので新鮮な感じですね。

それに伴って有名なアルゴリズムもいくつか実装したので、 ここではそういうアルゴリズムをMLではどう書いたのか紹介したいなと思います。

高速なべき乗のアルゴリズム

正式名称はよく分からないのですが、ありますよね?O(log n)でべき乗を計算するアルゴリズム。 x ^ nを計算する時、nが偶数の時は(x * x) ^ (n / 2)とみなすやつです。

OCamlで素朴に実装するとこんな感じ。

let rec power m = function
  | 0 -> 1.
  | n when n mod 2 = 0 ->
      power (m *. m) (n / 2)
  | n ->
      m *. power m (n - 1)

こんな計算も一瞬で終わります。

# (let n = 1000000000 in
power (1. +. 1. /. float_of_int n) n);;
- : float = 2.71828203081451081

ここでは浮動小数点数のべき乗を求める関数を記述しましたが、 単位元が存在して結合律が成り立つような二項演算子ならこのアルゴリズムでべき乗を計算できますよね?

そのような二項演算子のべき乗が必要になる度にいちいち関数の定義をやり直してはらちが明かないので、 高階関数として定義するのが僕の趣味です。

let rec power ( * ) one m = function
  | 0 -> one
  | n when n mod 2 = 0 ->
      power ( * ) one (m * m) (n / 2)
  | n ->
      m * power ( * ) one m (n - 1)

さっきの浮動小数点数のべき乗にも使えるし、

# (let n = 1000000000 in
power ( *. ) 1. (1. +. 1. /. float_of_int n) n);;
- : float = 2.71828203081451081

行列のべき乗にも使えますね。

# let fib n =
  (* O(log n)でフィボナッチ数を計算するけど、
     困ったことにO(n)で問題になるような入力ではオーバーフローする *)
  let
    ((a, _),
     (_, _)) =
    power
      (fun
        ((a, b),
         (c, d))
        ((a', b'),
         (c', d')) ->
           ((a * a' + b * c', a * b' + b * d'),
            (c * a' + d * c', c * b' + d * d')))
      ((1, 0),
       (0, 1))
      ((1, 1),
       (1, 0)) n in a;;
val fib : int -> int = <fun>
# Array.init 10 fib;;
- : int array = [|1; 1; 2; 3; 5; 8; 13; 21; 34; 55|]

べき乗の剰余を求めるのにも使ったりしてます。 abc034.contest.atcoder.jp

ダイクストラ

言わずと知れたやつです。

let rec dikjstra ( +. ) ( < ) q edge d =
  match q with
  | [] -> ()
  | u :: q ->
      let u, q =
        List.fold_left (fun (u, q') u' ->
          if d.(u) < d.(u') then (u, u' :: q')
          else (u', u :: q')) (u, []) q in
      List.iter (fun (v, c) ->
        if c +. d.(u) < d.(v) then
          d.(v) <- c +. d.(u)) (edge u);
      dikjstra ( +. ) ( < ) q edge d

ヒープを使っていない適当実装ですが、まぁ綺麗に書けたのではないかなと。 MLは手続き型言語として使っても良い言語です。

日本語版Wikipediaの図 と同じグラフの最短距離を求めたい場合はこんな感じに使います。

# let d = Array.make 6 infinity;;
val d : float array =
  [|infinity; infinity; infinity; infinity; infinity; infinity|]
# d.(0) <- 0.;;
- : unit = ()
# dijkstra ( +. ) ( < ) [0; 1; 2; 3; 4; 5] (function
    | 0 -> [ (1, 7.); (2, 9.); (5, 14.) ]
    | 1 -> [ (0, 7.); (2, 10.); (3, 15.) ]
    | 2 -> [ (0, 9.); (1, 10.); (3, 11.); (5, 2.) ]
    | 3 -> [ (1, 15.); (2, 11.); (4, 6.) ]
    | 4 -> [ (3, 6.); (5, 9.) ]
    | 5 -> [ (0, 14.); (2, 2.); (4, 9.) ]) d;;
      Warning 8: this pattern-matching is not exhaustive.
Here is an example of a value that is not matched:
6
- : unit = ()
# d;;
- : float array = [|0.; 7.; 9.; 20.; 20.; 11.|]

辺の重みが整数の場合は言わずもがな、 どの経路を通ったのか知りたい場合でも同じ関数を流用できます。 (先に通った経路からconsする都合上逆になってますが…)

# let d = Array.make 6 (infinity, []);;
val d : (float * '_a list) array =
  [|(infinity, []); (infinity, []); (infinity, []); (infinity, []);
    (infinity, []); (infinity, [])|]
# d.(0) <- (0., []);;
- : unit = ()
# dijkstra
    (fun (c, r) (c', rs) -> (c +. c', r :: rs))
    (fun (c, _) (c', _) -> c < c')
    [0; 1; 2; 3; 4; 5]
    (function
      | 0 -> [ (1, (7., "0->1")); (2, (9., "0->2")); (5, (14., "0->5")) ]
      | 1 -> [ (0, (7., "1->0")); (2, (10., "1->2")); (3, (15., "1->3")) ]
      | 2 -> [ (0, (9., "2->0")); (1, (10., "2->1")); (3, (11., "2->3"));
          (5, (2., "2->5")) ]
      | 3 -> [ (1, (15., "3->1")); (2, (11., "3->2")); (4, (6., "3->4")) ]
      | 4 -> [ (3, (6., "4->3")); (5, (9., "4->5")) ]
      | 5 -> [ (0, (14., "5->0")); (2, (2., "5->2")); (4, (9., "5->4")) ]) d;;
Warning 8: this pattern-matching is not exhaustive.
Here is an example of a value that is not matched:
6
- : unit = ()
# d;;
- : (float * string list) array =
[|(0., []); (7., ["0->1"]); (9., ["0->2"]); (20., ["2->3"; "0->2"]);
  (20., ["5->4"; "2->5"; "0->2"]); (11., ["2->5"; "0->2"])|]

ABC035のD問題でも同様の実装を使用したのですが、 ヒープを使わない適当実装が祟って部分点でした。 後でヒープも実装する必要がありそうですね。 abc035.contest.atcoder.jp

結び

今回の実装ではあまり関数型言語としてOCamlを使っていませんでしたが、 それでも汎用性と可読性に優れたコードになったのではないでしょうか。

書けば書くほどOCamlへの愛着が増していくようです。