【Common Lisp】
すべての組み合わせを、再帰とループで
列挙する
(可変長の多重ループのDFS)

  • N重ループの深さが実行時に決まるとき、forループだけでは入れ子の段数をコードに書けません
  • 再帰が「何段目か」という深さを担当し、loopが「その段でどの値を選ぶか」という幅を担当することで、可変長の多重ループが自然に書けます
  • 末端で何を返し、loopの蓄積節をどう選ぶかで、「数える」「集める」を切り替えられます
  • 再帰版では呼び出しスタックが状態を持ち、pushpopを使ったループ版ではリストが状態を持ちます

関連記事

1. 問題と全列挙

競技プログラミングの問題 を考えるときに、全列挙のコードが必要になりました。

N 個のサイコロを振る。
各サイコロの目は 1 以上 M 以下の整数。
出た目の合計が K 以下になる組み合わせの数を 998244353 で割った余りを求めよ。

素朴に考えれば、長さNの列を全列挙して、合計がK以下のものを数えればよいことになります。
もちろん、組み合わせはM^N通りなので、提出用には時間が足りません1
しかし、答えの性質を確かめて考え方を深める出発点として、素朴な解をスムーズに作れるように考え方を整理します。

1.1. なぜforループだけでは書きにくいのか

全列挙がややこしいのは、単純な for ループでは書きにくいからです。

もし、N = 3 に固定されているなら、3重ループで書けます。

(loop for a from 1 to M
      do (loop for b from 1 to M
               do (loop for c from 1 to M
                        do ...)))Code language: Lisp (lisp)

しかし、 N が実行時に決まる場合には、何重にすればよいかをコードに書けませんし、N = 50 なら50重ループが必要ですが、それを書くのは大変です。

1.2. 再帰とループを組み合わせた全探索

そこで、ループに再帰を組み合わせて、「残りのループ構造」を表せるようにします。

NMN^M の全探索(深さ優先探索)の基本形は、

(labels ((rec (i 蓄積引数...)
           (cond
             ((= i N)
              末端の処理)
             (t (loop for x from 1 to M
                      蓄積節 (rec (1+ i) 更新した引数...))))))
  (rec 0 初期値...))Code language: JavaScript (javascript)

具体的なコードで見てみます。

(defconstant +mod+ 998244353)

(defun count-dice-sum/naive (N M K)
  (labels ((rec (i sum)
             (cond
               ((= i N)
                (if (<= sum K) 1 0))
               (t
                (loop with count = 0
                      for x from 1 to M
                      do (incf count (rec (1+ i) (+ sum x)))
                      finally (return (mod count +mod+)))))))
    (rec 0 0)))Code language: Lisp (lisp)

内部の補助関数 reclabels2 で定義し、(rec i sum) で i 個目のサイコロを振り、これまでの合計が sum というように、状態を渡すようにします。

たとえば、2面のサイコロ(N = 3, M = 2)を3個振るなら、呼び出しはこう広がります。

(rec 0 0)
  x=1 -> (rec 1 1)
           x=1 -> (rec 2 2)
                    x=1 -> (rec 3 3)
                    x=2 -> (rec 3 4)
           x=2 -> (rec 2 3)
                    x=1 -> (rec 3 4)
                    x=2 -> (rec 3 5)
  x=2 -> (rec 1 2)
           x=1 -> (rec 2 3)
                    x=1 -> (rec 3 4)
                    x=2 -> (rec 3 5)
           x=2 -> (rec 2 4)
                    x=1 -> (rec 3 5)
                    x=2 -> (rec 3 6)

ループの x=..の部分で横に広がり、rec を呼ぶたびにもう1段深い loop が始まります。
末端の (rec 3 sum) がそれぞれ1つの列に対応します。
たとえば x=1, x=2, x=1 と選んだ経路は(rec 3 4)で合計 4 です。

これで、実行時には N 重ループのように振る舞うのです。
固定個数の多重ループはforループで手で書けますが、可変個数になると再帰で表すのが自然です3

2. 再帰の状態変数と基底ケース

(rec i sum)i は「何個まで出目を決めたか」、sum は「ここまでの合計」を表します。

再帰関数の isumのような引数を、「状態変数」といいます。
最初は、まだ何も選んでいないので (rec 0 0) から始めます。

再帰関数では、まず終了になる基底ケースを決めます。
今回は、i = N になったときで、「この1通りは有効か」を 1 または 0 で表します。

(cond
  ((= i N)
   (if (<= sum K) 1 0))
  (t
   (loop ...
         finally (return (mod count +mod+)))))Code language: Lisp (lisp)

出目の合計 sumK 以下なら 1 です。

2.1. recとloopの役割

loop for x from 1 to M は、その位置で選べる目を全部試します。

(loop for x from 1 to M
      sum (rec (1+ i) (+ sum x)) into count
      finally (return (mod count +mod+)))Code language: Lisp (lisp)

x に対して (rec (1+ i) (+ sum x)) を呼ぶことで、「1個追加したので深さを1進め、合計に x を加える」という意味になります。
つまり、再帰が「何段目か」という深さを、loop が「その段でどの値を選ぶか」という幅を担当しています。

再帰:位置を1つ進める(深さ)
loop:その位置でありうる値を全部試す(幅)

rec の返り値は「その x を選んだ場合に最終的に条件を満たす列の数」です。
それを count に積み上げて、最後に mod を取って返します4

2.2. 再帰の代わりに自分でスタックを管理する

全探索は、再帰を使わない場合には、リストを pushpop でスタックとして使っても書けます。

(defconstant +mod+ 998244353)

(defun count-dice-sum/stack (N M K)
  (loop with stack = (list (list 0 0))
        with count = 0
        while stack
        for (i sum) = (pop stack)
        if (= i N)
          do (incf count (if (<= sum K) 1 0))
        else
          do (loop for x from 1 to M
                    do (push (list (1+ i) (+ sum x))
                             stack))
        finally (return (mod count +mod+))))Code language: Lisp (lisp)

スタックには、次の処理を (i sum) というリストで積んで予約します。

  • 再帰版では、関数呼び出しのスタックが isum を保持しています。
    (rec 1 3) を呼ぶと、その呼び出しフレームが i=1, sum=3 という状態を覚えています。
  • ループ版では、自分で用意したリストが (1 3) という形で同じ状態を保持します。

最初は (list (list 0 0))、つまり「i=0, sum=0 という状態1つ」から始めています。
pop で取り出した状態を isum に分解し、再帰版と同じ条件で分岐します。
i = N なら判定、そうでなければ次の状態を全部 push して予約します5

つまり、再帰版では呼び出しスタックが (i sum) を持ち、ループ版ではリストスタックが (i sum) を持つという対応になっています。
つまり、ループ版では再帰が暗黙に管理していた「深さ」「戻り先」「終了条件」を、pushpoploop while stack として、自分で書く必要があるわけです。

3. 組み合わせの収集

count-dice-sum/naive は、末端で 1 または 0 を返し、loopincf して足し上げました。

この構造のまま、末端で返すものと loop の蓄積節を変えると、「数える」から「集める」に切り替えられます。

(defun collect-dice-sum/naive (N M K)
  (labels ((rec (i sum lst)
             (cond
               ((= i N)
                (if (<= sum K)
                    (list lst)
                    nil))
               (t
                (loop for x from 1 to M
                      append (rec (1+ i)
                                  (+ sum x)
                                  (cons x lst)))))))
    (rec 0 0 nil)))Code language: Lisp (lisp)

まず、状態変数に、lst を追加します。
ここには、そこまでのサイコロの出目をリストを追加していきます。
リストは、先頭に追加する方が速いので、逆順に並べていきます。

3.1. 基底ケースはリストのリストを返す

次に、基底ケースで返すものを考えます。

((= i N)
 (if (<= sum K)
     (list lst)
     nil))Code language: Lisp (lisp)

条件を満たす組み合わせそのものなので、追加してきた lst を返せばよさそうですが、注意が必要です。
loop で複数の結果をつなげるには、各結果がリストである必要があるからです。
そこで、有効なら (list lst)、無効なら nil を返します。
もし、正順で返したければ、(list (reverse lst)) とします。

3.2. 再帰ケースのループ

次に、loop の蓄積節を選びます。

(t
 (loop for x from 1 to M
       append (rec (1+ i)
                   (+ sum x)
                   (cons x lst))))Code language: Lisp (lisp)

リストをつなげるなら append が対応します6

lstcons で先頭に積むので逆順になります7

たとえば、N=2, M=3, K=4 で試すとこうなります。

(collect-dice-sum/naive 2 3 4)
;=> ((1 1) (2 1) (3 1) (1 2) (2 2) (1 3))Code language: PHP (php)

今回の問題は、順序は関係ありませんが、(2 1) は「先に1、次に2を選んだ」列なので、本来の順序では (1 2) です。

3.3. loopの蓄積節の選び方

count-dice-sum/naivecollect-dice-sum/naive を並べると、違いは3箇所だけです。

数える版:
  状態変数      (i sum)
  末端の返り値  (if (<= sum K) 1 0)
  loopの蓄積    sum (rec ...) into sum

集める版:
  状態変数     (i sum lst)
  末端の返り値  (if (<= sum K) (list lst) nil)
  loopの蓄積   append (rec ...)Code language: Lisp (lisp)

loop の蓄積節は、再帰の返り値の型に合わせて選びます。
返り値が数値なら sum で足し、返り値がリストなら append でつなげます。
条件で絞るなら collectwhen を組み合わせ、個数を数えるなら countwhen を組み合わせます。

構造は同じで、「末端で何を返すか」と「それをどう蓄積するか」の組み合わせが変わるだけです。

4. 動的計画法が必要になるとき

「再帰で深さ、ループで幅」という組み合わせは、「DFS(深さ優先探索)」として名前がついています。

「1つ選んで深く進み、末端まで到達したら戻って別の候補を試す」探索で8、全探索・組み合わせ生成・順列生成では、この形が繰り返し現れます。

「次の1つを選んで深く進む」操作を再帰に任せ、「その位置で選べる候補を全部試す」操作をループに任せることで、可変長の全探索が自然に書けます。

ただし、この素朴解には非効率な点があります。
同じ状態 (i sum) を複数の経路から何度も計算していることです。

たとえば (rec 2 3) は、(1, 2) と選んだ経路からも (2, 1) と選んだ経路からも現れます。
この重複を避けて表に保存し再利用するのが、DP(動的計画法)です9

  1. N = M = 50のとき、M^N = 50^50 ≈ 8.9 × 10^84通りになります。競技プログラミングでは一般に10^8程度の操作が1秒の目安とされるので、この規模の全列挙は現実的ではありません。
  2. labels はCommon Lispで局所関数を定義するスペシャルフォームです。flet との違いは、labels で定義した関数が自分自身を再帰的に呼べる点です。 – CLHS: Special Operator LABELS, FLET, MACROLET
  3. Common Lispは言語仕様として末尾再帰の最適化を保証していません。深い再帰はスタックオーバーフローを引き起こす可能性があります。この問題の制約ではN ≤ 50なので素朴解では問題になりませんが、深さが大きくなる場合はループ版やDP版への切り替えを検討する必要があります
  4. 998244353は競技プログラミングでよく使われる素数です。素数であるため剰余環における逆元(割り算)が扱いやすく、かつ32bit整数の上限(約21億)を超えないサイズなので、掛け算のオーバーフローを64bit整数で防ぎやすいという性質があります。 – 「998244353 で割ったあまり」の求め方を総特集! – Qiita
  5. pushpop はCommon Lispのマクロです。(push item place)(setf place (cons item place)) の短縮形、(pop place) は先頭要素を取り出して変数を更新します。関数ではなくマクロなので、変数を直接変更できます。 – CLHS: Macro PUSH
  6. append は最後の引数を除くすべてのリストをコピーして連結します。元のリストは変更されません。ただし、再帰のたびに append を呼ぶと結果リストのコピーが積み重なるため、大きなリストでは cons で先頭に積んでから最後に nreverse する方が効率的です。 – CLHS: Function APPEND
  7. (cons x lst)xlst の先頭に追加した新しいリストを返します。再帰の各段階で先頭に追加し続けるため、選んだ順序と逆向きに並びます。x=1, x=2 と選んだ場合、(cons 2 (cons 1 nil)) となり (2 1) になります。
  8. DFSはスタックを使って実装することもできます。再帰関数を使うと処理系の呼び出しスタックが暗黙のスタックとして機能するため、コードが簡潔になります。 – DFS(深さ優先探索)超入門! – Qiita
  9. このDice Sum問題のDP解では、dp[i][s] を「i個の出目を選んだときに合計がsになる組み合わせの数」として定義し、i を1から N まで更新していきます。計算量はO(N × K × M)になり、素朴解のO(M^N)から大幅に改善されます