amazing-blog

がんばります

AtCoder に登録したら解くべき精選過去問 10 問を OCaml で解いてみた

概要

OCaml競技プログラミングに取り組むことを推進する記事です。

qiita.com

で取り上げられている10問をOCamlで解きました。

OCamlで競プロ

OCamlは静的型付けの関数型言語です。
言語についてのアレコレはここで話しません

OCamlは以下のような点で競プロ向きだと思います。

  • 競プロするのに十分高速な入出力ライブラリがある(Scanf, Printfモジュール)
  • 結構速い
    • スクリプト言語ではTLE(時間制限超過)になってしまうような問題もなかなかの実行速度でやってくれてうれしい

逆に以下のような点では不利だと思います。

  • 標準ライブラリが充実していない
    • 動的配列とかPriority Queueとか欲しい…
      • (追記: Priority QueueはMap.min_bindingsで代用できそう)
    • AtCoderではBatteriesという便利なライブラリ群が使えるのでうれしい
  • 計算量を最小限におさえるために必要になったりする、ガシガシ副作用起こすプログラムを書くのは(手続き型言語に比べて)しんどい

本編

問題の一覧は↓

AtCoder Beginners Selection - AtCoder

1問目 : ABC086A - Product

方針

ab の積が偶数かどうかは a * b mod 2 = 0 で判定できます

実装

let () = Scanf.scanf "%d %d" (fun a b ->
  print_endline @@ if a*b mod 2 = 0 then "Even" else "Odd")

2問目 : ABC081A - Placing Marbles

方針

入力をどう取るかによって実装の仕方が変わってくると思います

  1. 3桁の数を1つの整数 a としたとき、各桁は以下のように取り出すことができます
    • 100の位 : a / 100
    • 10の位 : a / 10 mod 10
    • 1の位 : a mod 10
  2. Scanf.scanfを上手く使えば 3桁の数を1桁ずつ取ることもできます

実装

(* 実装 1 *)
let () = Scanf.scanf "%d" (fun a -> a/100 + a/10 mod 10 + a mod 10)
  |> Printf.printf "%d\n"
(* 実装 2 *)
let () = Scanf.scanf "%1d %1d %1d" (fun a b c -> a+b+c)
  |> Printf.printf "%d\n"

3問目 : ABC081B - Shift Only

方針

問題文のとおりにシミュレートしてもよいですが

  1. a_i について、「奇数にするには何回2で割る必要があるか」を求める( b_i とする )
  2. { b_i } の最小値を求める

というふうに解いてもよいので、この方針でやります

実装

let rec calc x = (* xを奇数にするために何回2で割る必要があるかを求めている *)
  if x mod 2 = 0 then 1 + calc (x / 2) else 0
let () =
  Scanf.scanf "%d" (fun n ->
    Array.init n (fun _ -> Scanf.scanf " %d" (fun x -> x)))
  |> Array.map calc (* 各要素にcalcを適用 *)
  |> Array.fold_left min 100 (* 100 は適当に大きな数 *)
  |> Printf.printf "%d\n"
  • 繰り返し入力を取る時は Array.init + Scanf.scanf が便利です
    • Listでもいいんですけど、現在のAtCoderで使えるOCamlのバージョンが古くてList.initがないため…
  • |> が多くて読みづらいかと思うかもしれませんが、これは書いてる人の趣味です…

4問目 : ABC087B - Coins

方針

硬貨の出し方は (A+1)(B+1)(C+1) 通りあるので、このそれぞれについて金額が X と一致するかを試します。数学っぽく言うと

{ (a, b, c) | 0 ≦ a ≦ A, 0 ≦ b ≦ B, 0 ≦ c ≦ C, 500×a + 100×b + 50×c = X }

の要素数を求めます

実装1

let rec range a b = if a > b then [] else a :: range (a+1) b
let (>>=) x f = List.map f x |> List.concat

let () =
  let a, b, c, x = Scanf.scanf "%d %d %d %d" (fun a b c x -> a, b, c, x) in
  (range 0 a >>= fun i ->
   range 0 b >>= fun j ->
   range 0 c >>= fun k -> [(i, j, k)])
  |> List.filter (fun (i, j, k) -> 500*i + 100*j + 50*k = x)
  |> List.length
  |> Printf.printf "%d\n"
  • ここでは List モナド っぽい書き方をしています

実装2

(2019/02/15 追記しました)

実装1 が分かりづらいというコメントを頂いたので、もうすこし愚直なやつを書きました

let solve a b c x =
  let rec f acc i j k =
    if i > a then acc
    else if j > b then f acc (i+1) 0     0
    else if k > c then f acc i     (j+1) 0
    else f (acc + if 500*i + 100*j + 50*k = x then 1 else 0) i j (k+1)
  in f 0 0 0 0

let () = Scanf.scanf "%d %d %d %d" solve |> Printf.printf "%d\n"

solveという関数の中のfという関数で、すべてのパターンを試しています。 引数accはこれまでに金額が一致したパターンの個数であり、引数i j kがこれから試すパターンにおけるそれぞれの硬貨の数です。 再帰的に呼び出されるfの引数i j kの数値を遷移を観察すると、

0 0 0
→ 0 0 1
...
→ 0 0 c
→ 0 0 (c+1) [skip]
→ 0 1 0
...
→ 0 (b+1) 0 [skip]
→ 1 0 0
...
→ (a+1) 0 0 [終了]

となっており、全てのパターンについて試していることが分かると思います。

実装1ではパターンのリストを作ってからfilterとかしていますが、実装2ではリストは生成されずループで処理されるので効率が良いです

5問目 : ABC083B - Some Sums

方針

各桁の和を求める関数を用意 → これを1~Nの整数に適用 → 和がA以上B以下のものを探す(情報量0)

実装

let rec range a b = if a > b then [] else a :: range (a+1) b
let rec digit_sum x = if x <= 0 then 0 else x mod 10 + digit_sum (x / 10)
let () =
  let n, a, b = Scanf.scanf "%d %d %d" (fun n a b -> n, a, b) in
  range 1 n
  |> List.filter (fun x -> a <= digit_sum x && digit_sum x <= b)
  |> List.fold_left (+) 0
  |> Printf.printf "%d\n"

6問目 : ABC088B - Card Game for Two

方針

2人のプレイヤーの最適戦略は、自分のターンで残っている最大得点のカードを取ることです。 この戦略はカードを得点の降順に並べて、前から交互にAliceとBobに与えることで再現できます。

このときの(Aliceの得点) - (Bobの得点) は

  a_0 - a_1 + a_2 - ... + (-1)^n * a_n
= (+ 1) * a_0  +  (- 1) * a_1  +  (+ 1) * a_2 ... + (-1)^n * a_n

実装

let rec list_init n f = if n <= 0 then [] else let x = f () in x :: list_init (n-1) f
let () =
  Scanf.scanf "%d" (fun n -> list_init n (fun _ -> Scanf.scanf " %d" (fun x -> x)))
  |> List.sort (-) |> List.rev
  |> List.fold_left (fun (a, s) e -> (a + s * e, -s)) (0, +1)
  |> fst |> Printf.printf "%d\n"

7問目 : ABC085B - Kagami Mochi

方針

大きい餅から小さい餅の順においていくのが最適です。問題の条件から、同じ大きさの餅は2つ置けません。 よって { d_i } の値の重複がなくなるように要素を削除した後の要素数が答えとなります。

実装 1

let rec list_init n f = if n <= 0 then [] else let x = f () in x :: list_init (n-1) f

let () =
  Scanf.scanf "%d" (fun n -> list_init n (fun _ -> Scanf.scanf " %d" (fun x -> x)))
  |> List.sort_uniq (-)
  |> List.length
  |> Printf.printf "%d\n"

List.sort_uniq は リストをソートする関数なのですが、それと同時に値の重複を除いてくれます。 これを適用した後のリストの要素数が答えになります

実装 2

List.sort_uniq を知らなかった人のものまね

let rec list_init n f = if n <= 0 then [] else let x = f () in x :: list_init (n-1) f

let solve xs =
  let rec f v = function
    | []       -> 0
    | (x::xs') -> if v = x then f w xs' else 1 + f x xs' in
  f 0 xs

let () =
  Scanf.scanf "%d" (fun n -> list_init n (fun _ -> Scanf.scanf " %d" (fun x -> x)))
  |> List.sort (-) |> solve
  |> Printf.printf "%d\n"

8問目 : ABC085C - Otoshidama

方針

10000円札がa枚で5000円札b枚だとすると、1000円札はN-a-b枚に決まります。 この時の金額の合計は 10000*a+5000*b+1000*(N-a-b) になるので、これがYになるようなペア(a, b)が存在するかを判定すればよいです

数学っぽく書くと

{ (a, b) | 0 ≦ a ≦ N, 0 ≦ b ≦ N, N-a-b ≧ 0, 10000×a + 5000×b + 1000×(N-a-b) = Y }

が空じゃなければOK

計算量的には O(N^2) になりますが、 N <= 2000 程度なのでOCamlなら間に合わせることができます

実装 1

let rec range a b = if a > b then [] else a :: range (a+1) b
let (>>=) x f = List.map f x |> List.concat

let () =
  Scanf.scanf "%d %d" (fun n y ->
  (range 0 n >>= fun a ->
   range 0 (n-a) >>= fun b -> [(a, b, n-a-b, 10000*a + 5000*b + 1000*(n-a-b))])
  |> List.filter (fun (_,_,_,p) -> p = y))
  |> (function
    | [] -> print_endline "-1 -1 -1"
    | ((a,b,c,_)::_) -> Printf.printf "%d %d %d\n" a b c)

長さ (N+1)*(N+2)/2 のリストを生成していて実行効率が微妙だったりします

実装 2

let rec solve a b n y =
  if a > n                              then (-1, -1, -1)
  else if a+b > n                       then solve (a+1) 0 n y
  else if 10*a + 5*b + (n-a-b) = y/1000 then (a,b,n-a-b)
  else solve a (b+1) n y

let () = let a, b, c = Scanf.scanf "%d %d" (solve 0 0) in Printf.printf "%d %d %d\n" a b c

再帰で答えを見つけるプログラムも用意しました。 こうすれば選択肢のリストは生成されず、答えが1つ見つかった段階で探索が打ち切りになるので効率が良いです

9問目 : ABC049C - 白昼夢 / Daydream

方針

文字列を先頭から見ていき判定することもできますが、末尾から見ていけばより簡単に判定することができます。

実装 1

let list_of_str s = let rec f i = if i >= String.length s then [] else s.[i] :: f (i+1) in f 0
let rec solve = function
  | [] -> "YES"
  | 'm'::'a'::'e'::'r'::'d'::xs | 'r'::'e'::'m'::'a'::'e'::'r'::'d'::xs
  | 'e'::'s'::'a'::'r'::'e'::xs | 'r'::'e'::'s'::'a'::'r'::'e'::xs -> solve xs
  | _ -> "NO"
let () = Scanf.scanf "%s" list_of_str |> List.rev |> solve |> print_endline

末尾から判定しています

実装 2

(2019/02/15 実装変えました)

let list_of_str s =  let rec f i = if i >= String.length s then [] else s.[i] :: f (i+1) in f 0
let rec solve = function
  | [] -> "YES"
  (* 1. erase[r] はその場で消費して OK *)
  | 'e'::'r'::'a'::'s'::'e'::'r'::xs
  | 'e'::'r'::'a'::'s'::'e'::xs
  (* 2. dreamer... と来ても実は dream erase[r] かもしれない*)
  | 'd'::'r'::'e'::'a'::'m'::'e'::'r'::'a'::'s'::'e'::'r'::xs
  | 'd'::'r'::'e'::'a'::'m'::'e'::'r'::'a'::'s'::'e'::xs
  (* 3. dream[er] *)
  | 'd'::'r'::'e'::'a'::'m'::'e'::'r'::xs
  | 'd'::'r'::'e'::'a'::'m'::xs -> solve xs
  (* 4. NO *)
  | _ -> "NO"

let () = Scanf.scanf "%s" list_of_str |> solve |> print_endline

先頭から判定する場合は、末尾から判定する場合よりも少しだけ工夫が必要になります。 erase eraser が先頭にあった場合はその場で消費してOKです。 ところが dreamer が先頭に来たとき、この中の erdream と合わせるべきかどうかを決定するには更に先読みを行う必要があります。 具体的には dreamerase... という文字列があった時に、 dreamer ase... とせずに dream erase... とみなす必要があります。

10問目 : ABC086C - Traveling

方針

時刻 t_i(x_i, y_i) を出発し、 t_{i+1} 秒後に (x_{i+1}, y_{i+1}) に存在できる条件は以下の2つです

  1. 距離 |x_i-x_{i+1}| + |y_i-y_{i+1}| が、移動できる時間 t_{i+1}-t_i 以下
  2. 余った時間 t_{i+1} - t_i - |x_i-x_{i+1}| - |y_i-y_{i+1}| でちょうど (x_{i+1}, y_{i+1}) に到着

これが全ての隣接する2つの地点の間で成立していればOK

実装

let () =
  Scanf.scanf "%d" (fun n -> Array.init n (fun _ -> Scanf.scanf " %d %d %d" (fun t x y -> t, x, y)))
  |> Array.fold_left (fun (b,(t1,x1,y1)) (t2,x2,y2) ->
      let d = abs (x1-x2) + abs (y1-y2) in
      let b' = b && t2-t1 >= d && (t2-t1-d) mod 2 = 0 in
      (b',(t2,x2,y2))) (true,(0,0,0))
  |> (fun (b,_) -> if b then "Yes" else "No")
  |> print_endline

まとめ

  • OCamlは速度の割に簡潔に書ける(例外あり)ので、そこそこ競プロ向きだと思います
  • AtCoderで使えるOCamlのバージョン上げて欲しい Batteries使えるからそんな困らない気がしてきた