Schemeで桁DPに入門してみた

はじめに

こちらの記事

のScheme版を書いている。途中、桁DPに関して詳しいこちらの記事が紹介されていた。

桁DPは数の性質を利用して、「条件に合う数」を「条件に合う状態遷移」ととらえる独特の考え方を必要とする。理解のために元記事に書かれている実装をSchemeで再実装してみたのだが、長くなってしまったので別記事としてまとめてみた。

基本的には元記事に書かれている実装をSchemeで再実装したものなので、コードが理解できない場合は元記事を適宜参照してみてください。

A以下の非負整数の総数を求める

(use gauche.array)
(use util.combinations)

; 数字を文字リストに変換する
(define (number->list n)
  (string->list (number->string n)))  
; 数を1桁数字リストに変換する
(define (integer->intlist n)
  (map (^(c) (- (char->integer c) 48)) (number->list n)))

; --------------------------------------------------------------------------------

; 状態の定数
(define *exact* 0)
(define *less* 1)

; A以下の非負整数の総数を求める
(define (count-number end)
  (let* [(max-digit (string-length (number->string end)))
         (dp (make-array (shape 0 (+ max-digit 1) 0 2) 0))
         (a (list->vector (integer->intlist end))) ]
    ; 初期条件第0桁のコストはexact状態のみ1で確定
    (array-set! dp 0 *exact* 1)
    ; 計算本体
    (map (^(x) (let-values [((i less?) (apply values x))]
                 ; 次の桁が自由（0~9を取りうる）か、Aに制約されるか
                 (let* [(max-d (if (= less? *less*) 9 (vector-ref a i)))]
                   ; 次の桁への遷移条件フラグの設定
                   (map (^(d) (let [(less?_ (if (or (= less? *less*) (< d max-d)) *less* *exact*))]
                                (array-set! dp (+ i 1) less?_
                                            (+ (array-ref dp (+ i 1) less?_) (array-ref dp i less?)))))
                        ; 取りうる桁の数字を全て列挙する
                        (iota (+ max-d 1))))))
         ; 遷移の各状態を全て列挙する
         (cartesian-product (list (iota max-digit) '(0 1))))
    ; 結果
    (+ (array-ref dp max-digit *less*)
       (array-ref dp max-digit *exact*))))

(count-number 9023)
;  => 9024

桁DPの考え方を忘れてしまったときのためにメモ。

• 桁DPは与えられた数Aに対して、0以上A以下のすべての数を数え上げる。
• 数え上げは、大きい位から小さい位に向かって順に状態を確定させていくことで行う。次の位の桁に遷移するときにどのような状態（つまり0～9のどの数字か）を取りうるかを計算し、次の桁の状態（＝大体は大きい位の状態数ｘその位で取りうる状態数）を確定させていく。
• 最も基本の「状態」は、数Aに対してぴったり一致する（exact）か数Aよりも小さい(less)かの2択。判定条件が増えるたびに、状態の数も増やしていく。
• 最初に数Aの桁数（max-digit）を計る。次にその桁数+1のDP用のキャッシュ（この実装ではarray）を用意する。桁数+1なのは初期状態として第0桁が必要だから。キャッシュの次元は桁数x状態数分だけ確保できるようにする。
• 第0桁から始めて最後のmax-digit桁（Aにおける1の位）の状態まで確定させたときの、求める状態のmax-digit桁の取りうる状態数が答えとなる。

実装における留意点としては、

• 元記事ではDP用のキャッシュにdefaultdictを使っていたがarrayで実装している。
• 桁iを考えるとき、i=0からmax-digitまでループでカウントアップしてしまいがちだが、(桁数 状態A 状態B 状態C ...)という状態セットを全パターン先に作ってから順番に評価していったほうが、見通しが良く応用しやすいコードになる。
• 計算本体の中では、hoge?をi桁での状態、hoge?_をi+1桁の状態としている。
• 初期状態は(0桁 0 0 0 ... 0)で始まるようにしてあるため、この初期状態の状態数を1にセットしておく。

A以下の非負整数のうち、3の付く数の総数を求める

; 状態の定数
(define *not-has3* 0)
(define *has3* 1)

;A以下の非負整数のうち、3の付く数の総数を求める
(define (count-has3-number end)
  (let* [(max-digit (string-length (number->string end)))
         (dp (make-array (shape 0 (+ max-digit 1) 0 2 0 2) 0))
         (a (list->vector (integer->intlist end))) ]
    ; 初期条件第0桁のコストはexact状態のみ1で確定
    (array-set! dp 0 *exact* *not-has3* 1)
    ; 計算本体
    (map (^(x) (let-values [((i less? has3?) (values (car x) (cadr x) (caddr x)))]
                 (let* [(max-d (if (= less? *less*) 9 (vector-ref a i)))]
                   (map (^(d) (let [(less?_ (if (or (= less? *less*) (< d max-d)) *less* *exact*))
                                    (has3?_ (if (or (= has3? *has3*) (= d 3)) *has3* *not-has3*))]
                                (array-set! dp (+ i 1) less?_ has3?_
                                            (+ (array-ref dp (+ i 1) less?_ has3?_) (array-ref dp i less? has3?)))))
                        (iota (+ max-d 1))))))
         (cartesian-product (list (iota max-digit) '(0 1) '(0 1))))
    (+ (array-ref dp max-digit *less* *has3*)
       (array-ref dp max-digit *exact* *has3*))))

(count-has3-number 15)
;  => 2

3のつく数かどうか、を判定しなければならないので状態（has3?）が1つ増えた。最初に判定する状態をすべて列挙しているので、各状態で桁の数が3かどうかをチェックするコードを追加するだけで良い。

A以下の非負整数のうち、「3が付くまたは3の倍数」の数の総数を求める

は飛ばします。

A以下の非負整数のうち、「3が付くまたは3の倍数」かつ「8の倍数でない」数の総数を求める

(define *not-has3* 0)
(define *has3* 1)

;A以下の非負整数のうち、「3が付くまたは3の倍数」かつ「8の倍数でない」数の総数を求める
(define (count-has3-no8mult-number end)
  (let* [(max-digit (string-length (number->string end)))
         (dp (make-array (shape 0 (+ max-digit 1) 0 2 0 2 0 3 0 8) 0))
         (a (list->vector (integer->intlist end))) ]
    ; 初期条件第0桁のコストはexact状態のみ1で確定
    (array-set! dp 0 *exact* *not-has3* 0 0 1)
    ; 計算本体
    (map (^(x) (let-values [((i less? has3? mod3 mod8) (apply values x))]
                 (let* [(max-d (if (= less? *less*) 9 (vector-ref a i)))]
                   (map (^(d) (let [(less?_ (if (or (= less? *less*) (< d max-d)) *less* *exact*))
                                    (has3?_ (if (or (= has3? *has3*) (= d 3)) *has3* *not-has3*))
                                    (mod3_ (modulo (+ mod3 d) 3))
                                    (mod8_ (modulo (+ (* mod8 10) d) 8))]
                                (array-set! dp (+ i 1) less?_ has3?_ mod3_ mod8_
                                            (+ (array-ref dp (+ i 1) less?_ has3?_ mod3_ mod8_)
                                               (array-ref dp i less? has3? mod3 mod8)))))
                        (iota (+ max-d 1))))))
         (cartesian-product (list (iota max-digit) '(0 1) '(0 1) (iota 3) (iota 8))))
    ; 結果
    (apply +
           (map (^(x) (let-values [((less? has3? mod3 mod8) (apply values x))]
                        (array-ref dp max-digit less? has3? mod3 mod8)))
                (filter-map (^(x) (if (and (or (= (cadr x) *has3*) (= (caddr x) 0)) (not (zero? (cadddr x))))
                                    x
                                    #f)) 
                            (cartesian-product (list '(0 1) '(0 1) (iota 3) (iota 8))))))))

(count-has3-no8mult-number 24)
;  => 9

24までには3の倍数は8個あるが、13があるので+1、24は8の倍数なので-1して結果9になっている。

modulo計算をしており、mod 3の結果とmod 8の結果を状態として持つため、arrayは次元を増やしている。状態数はmodの結果である0～2、0～7を取れるようにしている。

結果の出力が複雑になっているが、最初に全状態を列挙し条件にそぐわないものをふるい落としている。

A以下B以上の非負整数のうち、「3が付くまたは3の倍数」かつ「8の倍数でない」数の総数を求める

も飛ばします。

終わりに

いちいち数え上げるのはまだるっこしく感じるが、全状態を最初に列挙してしまえば、あとは実にシンプルなんだな……。