fetburner.core

コアダンプ

OCamlで色々列挙する

AtCoder Beginner ContestのC問題で全探索が出題されがちな今日この頃,いかがお過ごしでしょうか? C++で競プロをやってる人達は重複順列に対して全探索して解を得たい時なんかに,いちいち重複順列を列挙する処理も含んだ再帰関数を書き直してるらしいですが,我々からすれば重複順列を列挙する処理を切り出して使い回せるようにしたいと考えるのが自然ではないでしょうか.

本記事では順列,組み合わせ,重複順列,重複組合せについて,十分速い時間で全列挙できる関数をOCamlで実装します.

重複順列

まずは一番簡単そうな,重複順列を列挙する関数を実装しようと思います.

こういうプログラムを実装する時,Haskellの人はリストモナドを使うんじゃないですかね. 例えば,xsの要素から重複を許してn要素を選ぶ順列をリストとして列挙する関数reperm*1は,以下のように書けます.*2

reperm 0 xs = [[]]
reperm n xs = do
  x <- xs
  zs <- reperm (n - 1) xs
  return (x : zs)

最近のOCamlならこれに近い書き方で同様に実装できるのですが,より高速な実装を得るためにdo記法を脱糖してプログラム変換をしていきましょう.

reperm 0 xs = [[]]
reperm n xs =
  xs >>= \x ->
    reperm (n - 1) xs >>= \zs ->
      return (x : zs)

fmap f xs = xs >>= return . fが成り立つので,reperm (n - 1) xs >>= \zs -> return (x : zs)の部分はfmapで十分ですね.

reperm 0 xs = [[]]
reperm n xs =
  xs >>= \x ->
    fmap (x : ) (reperm (n - 1) xs)

リストに対してのbind>>=(xs >>= f) = concat (map f xs)で定義されていたのを思い出すと*3,上のHaskellのコードはこんな感じのOCamlのコードに直せそうです.

let rec reperm n xs =
  if n <= 0
  then [[]]
  else List.concat @@
       List.map (fun x ->
         List.map (List.cons x) @@
         reperm (n - 1) xs) xs

ここで,List.concat = Fun.flip (List.fold_right ( @ )) []が成り立つので,List.fold_rightList.mapを融合変換してしまいましょう.

let rec reperm n xs =
  if n <= 0
  then [[]]
  else List.fold_right (fun x acc ->
         List.map (List.cons x) (reperm (n - 1) xs) @ acc) xs []

n-1要素を選ぶ重複順列を列挙してから,選んでおいたxを後からList.mapで追加する辺りが非効率的っぽいです.アキュムレーターを増やして何とかしたいですね.

reperm_aux ys n xs = List.map (( @ ) (List.rev ys)) (reperm n xs)となる関数reperm_auxを用いてrepermを書き直しましょう.

let reperm_aux ys n xs = List.map (( @ ) (List.rev ys)) (reperm n xs)

reperm n xsの定義を展開します.

let reperm_aux ys n xs =
  List.map (( @ ) (List.rev ys)) @@
  if n <= 0
  then [[]]
  else List.fold_right (fun x acc ->
         List.map (List.cons x) (reperm (n - 1) xs) @ acc) xs []

ifのthen節とelse節にList.map (( @ ) (List.rev ys))を分配します.

let reperm_aux ys n xs =
  if n <= 0
  then List.map (( @ ) (List.rev ys)) [[]]
  else List.map (( @ ) (List.rev ys)) @@
       List.fold_right (fun x acc ->
         List.map (List.cons x) (reperm (n - 1) xs) @ acc) xs []

List.map (( @ ) (List.rev ys)) [[]][List.rev ys]に簡約できますよね.

let reperm_aux ys n xs =
  if n <= 0
  then [List.rev ys]
  else List.map (( @ ) (List.rev ys)) @@
       List.fold_right (fun x acc ->
         List.map (List.cons x) (reperm (n - 1) xs) @ acc) xs []

ここで,xs = [x1; x2; ... xn]とすると,

List.fold_right (fun x acc ->
  List.map (List.cons x) (reperm (n - 1) xs) @ acc) xs []

の部分は

List.map (List.cons x1) (reperm (n - 1) xs) @
List.map (List.cons x2) (reperm (n - 1) xs) @
...
List.map (List.cons xn) (reperm (n - 1) xs)

を計算していると考えられるので,List.map f (l @ l') = List.map f l @ List.map f l'で繰り返し書き換えると,

List.map (( @ ) (List.rev ys))
  List.fold_right (fun x acc ->
    List.map (List.cons x) (reperm (n - 1) xs) @ acc) xs []
= List.map (( @ ) (List.rev ys))
   (List.map (List.cons x1) (reperm (n - 1) xs) @
    List.map (List.cons x2) (reperm (n - 1) xs) @
    ...
    List.map (List.cons xn) (reperm (n - 1) xs))
= List.map (( @ ) (List.rev ys))
    (List.map (List.cons x1) (reperm (n - 1) xs)) @
  List.map (( @ ) (List.rev ys))
    (List.map (List.cons x2) (reperm (n - 1) xs)) @
  ...
  List.map (( @ ) (List.rev ys))
    (List.map (List.cons xn) (reperm (n - 1) xs))
= List.fold_right (fun x acc ->
    List.map (( @ ) (List.rev ys))
      (List.map (List.cons x) (reperm (n - 1) xs))
    @ acc) xs []

よって

let reperm_aux ys n xs =
  if n <= 0
  then [List.rev ys]
  else List.fold_right (fun x acc ->
         List.map (( @ ) (List.rev ys))
           (List.map (List.cons x) (reperm (n - 1) xs))
         @ acc) xs []

List.map f (List.map g xs)みたいな形の式を見ると当然List.map (fun x -> f (g x)) xsに融合変換したくなってくるので

let reperm_aux ys n xs =
  if n <= 0
  then [List.rev ys]
  else List.fold_right (fun x acc ->
         List.map (( @ ) (List.rev (x :: ys))) (reperm (n - 1) xs)
         @ acc) xs []

ちょっと恣意的な書き換えですね.List.rev (x :: ys) @ zs = List.rev ys @ (x :: zs)なので,( @ ) (List.rev ys)List.cons xを関数合成すると( @ ) (List.rev (x :: ys))になります.

ここで,reperm_aux ys n xs = List.map (( @ ) (List.rev ys)) (reperm n xs)が成り立つようにreperm_auxを定義していたことを思い出すと,

let rec reperm_aux ys n xs =
  if n <= 0
  then [List.rev ys]
  else List.fold_right (fun x acc ->
         reperm_aux (x :: ys) (n - 1) xs @ acc) xs []

と書き換えられ,はれてList.mapを葬り去ることができました.

そうなるとリストの結合( @ )も消し去りたくなりますよね. 次はreperm_aux' ys n xs acc = reperm_aux ys n xs @ accとなる関数reperm_aux'を考えますが,これは明らかに

let rec reperm_aux' ys n xs acc =
  if n <= 0
  then List.rev ys :: acc
  else List.fold_right (fun x ->
         reperm_aux' (x :: ys) (n - 1) xs) xs acc

で定義できます. ちなみにList.fold_rightはスタックオーバーフローの危険があって使わない方が良いんですが,List.fold_rightだとスタックが溢れるほどデカい入力に対応したところで今度はヒープが溢れるだけなので考えないことにします.

あとはreperm_aux'を用いてrepermを定義し直せば完成なんですが…List.rev ysの所でシェアリングがぶっ壊れててメモリ使用量が増えてるんですよね.ちょっと仕様を変えてしまって,重複順列を逆にしたものが出てくるようにしましょう.

let rec reperm_aux' ys n xs acc =
  if n <= 0
  then ys :: acc
  else List.fold_right (fun x ->
         reperm_aux' (x :: ys) (n - 1) xs) xs acc
let reperm n xs = reperm_aux' [] n xs []

xsの要素数mとすると,時間計算量と空間計算量共にO(m^n)です.せやろな.

実際に使ってみるとこんな感じです.

# reperm 3 [1; 2];;
- : int list list =
[[1; 1; 1]; [2; 1; 1]; [1; 2; 1]; [2; 2; 1]; [1; 1; 2]; [2; 1; 2]; [1; 2; 2];
 [2; 2; 2]]
# reperm 2 [1; 2; 3];;
- : int list list =
[[1; 1]; [2; 1]; [3; 1]; [1; 2]; [2; 2]; [3; 2]; [1; 3]; [2; 3]; [3; 3]]

組み合わせ

後述するように,実は順列を列挙するより組み合わせを列挙する方が簡単です.組み合わせを列挙するプログラムを少しアレンジすれば重複組合せを列挙するプログラムも手に入るので,次は組み合わせの列挙に着手しましょう.

まず,効率は抜きにして直感的に正しそうなプログラムを書いてみましょう.リストxsからn要素を選ぶ組み合わせを列挙する関数comb n xsは,次のように書けます.

let rec comb n xs =
  match n, xs with
  | 0, _ -> [[]]
  | _, [] -> []
  | n, x :: xs -> List.map (List.cons x) (comb (n - 1) xs) @ comb n xs

0要素を選ぶ組み合わせは何も選ばない一択なので,comb 0 xs = [[]]ですね.[[]] <> []であることに注意してください.

何もないところから1個以上の要素を選ぶことはできないので,n>0のときcomb n [] = []で定義しています.

パスカルの三角形の直感に従えば分かりやすいと思うんですが,x :: xsの中からn要素を選ぶ組み合わせは, xを選ぶことを確定させた上でn-1要素をxsの中から選ぶ組み合わせと,xを選ばずにxsの中からn要素を選ぶ組み合わせの二つに分けられます. なので再帰呼び出しでその両者を列挙し,n>0のときのcomb n (x :: xs)を与えています.

重複順列のときと同じように,効率的なプログラムに直していきましょう.まずはList.mapを消すためにアキュムレーターを導入しましょう. comb_aux ys n xs = List.map (( @ ) (List.rev ys)) (comb n xs)となる関数comb_aux ys n xsは,以下のように定義できます.

let rec comb_aux ys n xs =
  match n, xs with
  | 0, _ -> [List.rev ys]
  | _, [] -> []
  | n, x :: xs -> comb_aux (x :: ys) (n - 1) xs @ comb_aux ys n xs

導出の際の議論も重複順列のときと同じような感じです.

次に,( @ )を消すためにアキュムレーターを導入します.comb_aux' ys n xs acc = comb_aux ys n xs @ accとなる関数comb_aux'は, 以下のように定義できます.

let rec comb_aux' ys n xs acc =
  match n, xs with
  | 0, _ -> List.rev ys :: acc
  | _, [] -> acc
  | n, x :: xs -> comb_aux' (x :: ys) (n - 1) xs (comb_aux' ys n xs acc)

あとはcomb_aux'を使ってcombを再定義するだけなんですが,例のごとくList.rev ysがウザいので逆順のまま列挙することにします.

let rec comb_aux' ys n xs acc =
  match n, xs with
  | 0, _ -> ys :: acc
  | _, [] -> acc
  | n, x :: xs -> comb_aux' (x :: ys) (n - 1) xs (comb_aux' ys n xs acc)
let comb n xs = comb_aux' [] n xs []

xsの要素数mとすると,時間計算量,空間計算量共にO({}_m \mathrm{C}_n)です.それはそう.

試しに動かしてみましょう.

# comb 2 [1; 2; 3];;
- : int list list = [[2; 1]; [3; 1]; [3; 2]]
# comb 3 [1; 2; 3; 4; 5];;
- : int list list =
[[3; 2; 1]; [4; 2; 1]; [5; 2; 1]; [4; 3; 1]; [5; 3; 1]; [5; 4; 1]; [4; 3; 2];
 [5; 3; 2]; [5; 4; 2]; [5; 4; 3]]

心の目で各要素にList.revを適用してやると,ちゃんと組み合わせが列挙できてそうな気がしますよね.

重複組合せ

実のところ,組み合わせを列挙するプログラムを少し修正するだけで重複組合せを列挙できます. リストx :: xsからxを選んだ上で残りのn-1要素を選ぶ際,重複組合せではまたxを選び直しても良くなるので,

| n, x :: xs -> comb_aux' (x :: ys) (n - 1) xs (comb_aux' ys n xs acc)

の部分を

| n, x :: xs' -> repcomb_aux (x :: ys) (n - 1) xs (repcomb_aux ys n xs' acc)

に変えるだけで重複組合せ列挙の補助関数repcomb_aux ys n xs accが完成します.

従って,リストxsから重複を許してn要素を選ぶ組合せを列挙する関数repcomb n xsは次のように定義できます.

let rec repcomb_aux ys n xs acc =
  match n, xs with
  | 0, _ -> ys :: acc
  | _, [] -> acc
  | n, x :: xs' -> repcomb_aux (x :: ys) (n - 1) xs (repcomb_aux ys n xs' acc)
let repcomb n xs = repcomb_aux [] n xs []

試しに動かしてみるとこんな感じです.

# repcomb 2 [1; 2; 3];;
- : int list list = [[1; 1]; [2; 1]; [3; 1]; [2; 2]; [3; 2]; [3; 3]]
# repcomb 2 [1; 2];;
- : int list list = [[1; 1]; [2; 1]; [2; 2]]

相変わらず各要素が逆になっているのはご愛嬌.

順列

簡単そうに見えてこいつが一番列挙しづらいです。

リストxsからn個の要素を選ぶ順列を列挙する関数perm n xsを実装しようと思うと,xsの各要素xについて, xsからxだけを取り除いたリストからn-1個の要素を選ぶ順列を求めることになると思います. イメージとしてはこんな感じ.

let rec select trace = function
  | [] -> [] 
  | x :: xs -> (List.rev trace, x, xs) :: select (x :: trace) xs
(* xsの各要素xと,xより左にある要素のリスト,xより右にある要素のリストを列挙する *)
let select xs = select [] xs

(* 与えられたリストからn要素を選ぶ順列 *)
let rec perm n xs =
  match n, xs with
  | 0, _ -> [[]]
  | _, [] -> []
  | n, xs ->
      List.concat @@
      List.map (fun (ys, x, zs) ->
        List.map (List.cons x) @@
        (* リストxsからxだけを取り除いたリストは,ys @ zs *)
        perm (n - 1) (ys @ zs)) @@
      select xs

ここでのxsからxだけを取り除いたリストを作る操作ys @ zsが鬼門で,純粋関数型に書こうとするとxより左にある要素の数(すなわちList.length ys)に比例した時間がかかってしまいます. これを各xについてやる訳ですから,再帰呼び出しごとにxsの要素数の二乗に比例した計算時間が必要になってしまいますね.これはまずい…

幸いOCaml手続き型言語ですから,関数プログラミングさえ諦めてしまえばリストxsから要素xの削除を定数時間で行うことができます. 破壊的代入によって高速化したperm n xsの実装を以下に示します.

type 'a mutable_list = Nil | Cons of 'a * 'a mutable_list ref

let perm n xs =
  let head = ref Nil in
  (* xsを変更可能なリストに変換 *)
  ignore (List.fold_left (fun tail x ->
    let tail' = ref Nil in
    tail := Cons (x, tail'); tail') head xs);
  let rec perm_aux ys n ptr acc =
    if n <= 0
    then ys :: acc
    else match !ptr with
         | Nil -> acc
         | (Cons (x, next)) as here ->
             let acc = perm_aux ys n next acc in
             ptr := !next; (* リストからconsセルhereを削除 *)
             (* consセルを削除したリストから,n-1要素を選ぶ *)
             let acc = perm_aux (x :: ys) (n - 1) head acc in
             ptr := here; (* consセルを戻す *)
             acc in
  perm_aux [] n head []

なんかもう凄いことになってますね.好き放題やった甲斐あって,xsの要素数mとすると時間計算量空間計算量共にO({}_m \mathrm{P}_n)です.

リストをいちいち変更可能なリストに変換してる辺りが無駄に見えるんですが, 黒魔術を使って普通のリストのconsセルもcdrを破壊的に書き換えられたりしませんか?邪悪にも程があるのでやりたくないですが.

試しにこの関数perm n xsを使ってみるとこんな感じです.

# perm 2 [1; 2; 3];;
- : int list list = [[2; 1]; [3; 1]; [1; 2]; [3; 2]; [1; 3]; [2; 3]]
# perm 3 [1; 2; 3];;
- : int list list =
[[3; 2; 1]; [2; 3; 1]; [3; 1; 2]; [1; 3; 2]; [2; 1; 3]; [1; 2; 3]]

一応辞書順にはなってるんですが,順列を反転したものが出てくるのはご愛嬌です.シェアリング解きたくないですし.

まとめ

順列,組み合わせ,重複順列,重複組合せについて,十分速い時間で全列挙できる関数をOCamlで実装しました. 最終的に得られた各関数の実装を以下に示します.

let rec reperm_aux' ys n xs acc =
  if n <= 0
  then ys :: acc
  else List.fold_right (fun x ->
         reperm_aux' (x :: ys) (n - 1) xs) xs acc
let reperm n xs = reperm_aux' [] n xs []

let rec comb_aux' ys n xs acc =
  match n, xs with
  | 0, _ -> ys :: acc
  | _, [] -> acc
  | n, x :: xs -> comb_aux' (x :: ys) (n - 1) xs (comb_aux' ys n xs acc)
let comb n xs = comb_aux' [] n xs []

let rec repcomb_aux ys n xs acc =
  match n, xs with
  | 0, _ -> ys :: acc
  | _, [] -> acc
  | n, x :: xs' -> repcomb_aux (x :: ys) (n - 1) xs (repcomb_aux ys n xs' acc)
let repcomb n xs = repcomb_aux [] n xs []

type 'a mutable_list = Nil | Cons of 'a * 'a mutable_list ref

let perm n xs =
  let head = ref Nil in
  (* xsを変更可能なリストに変換 *)
  ignore (List.fold_left (fun tail x ->
    let tail' = ref Nil in
    tail := Cons (x, tail'); tail') head xs);
  let rec perm_aux ys n ptr acc =
    if n <= 0
    then ys :: acc
    else match !ptr with
         | Nil -> acc
         | (Cons (x, next)) as here ->
             let acc = perm_aux ys n next acc in
             ptr := !next; (* リストからconsセルhereを削除 *)
             (* consセルを削除したリストから,n-1要素を選ぶ *)
             let acc = perm_aux (x :: ys) (n - 1) head acc in
             ptr := here; (* consセルを戻す *)
             acc in
  perm_aux [] n head []

*1:もっと良い関数名は無かったのか

*2:説明のために雰囲気で書いてるので,間違ってたらごめんなさい

*3:正確にはconcatMapだったかも