amazing-blog

がんばります

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

概要

qiita.com

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

OCamlで競プロ

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

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

  • 競プロするのに十分高速な入出力ライブラリがある(Scanf, Printfモジュール)
  • 結構速い

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

  • 標準ライブラリが充実していない
    • 動的配列とかPriority Queueとか欲しい…
    • AtCoderではBatteriesという便利なライブラリ群が使えるのでうれしい
  • ガシガシ副作用起こすプログラムを書くのは(手続き型言語に比べて)しんどい

筆者のOCaml歴あんまり長くないので拙いコードでもゆるして

本編

問題の一覧は↓

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 }

の要素数を求めます

実装

(* range a b = [a; a+1; ... ; b] *)
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"
  • OCamlにも for ループがあるので、3つの変数a b cについてループを回すこともできます
  • ここでは List モナド っぽい書き方をしています

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) (* a以上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

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 erase = function
  | ('e'::'r'::'a'::'s'::'e'::'r'::xs)
  | ('e'::'r'::'a'::'s'::'e'::xs)      -> Some xs
  | _ -> None

let rec solve xs =
  match erase xs with
  | Some xs' -> solve xs'
  | None -> match xs with
    | []                                       -> "YES"
    | ('d'::'r'::'e'::'a'::'m'::'e'::'r'::xs')
      when erase ('e'::'r'::xs') = None        -> solve xs'
    | ('d'::'r'::'e'::'a'::'m'::xs')           -> solve xs'
    | _ -> "NO"

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

先頭から判定しています(これより簡潔に書けるかもしれない)

バックトラック法でもいいと思います

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のバージョン上げて欲しい、-O2オプションもつけてよいかもしれない