🦄

桁DPとleading zero

2022/01/05に公開

競技プログラミングにおけるいわゆる「桁DP」において、leading zeroを考慮しなければならない場合があります。本記事では桁DPにおけるleading zeroの扱いについて解説します。

桁DPの基本的な事柄については、例えば以下の記事を参照してください。

解説

次の問題を考えてみましょう。(制約は1 \leq N \leq 10^{1000}, 1 \leq K \leq 100とします。実際には10^9 + 7などで割った余りを答えるケースが多いと思いますが、以下の例では省略します。)

問題1

1以上N以下の整数で、各桁の数字の和がK以下のものはいくつあるか?

これは桁DPの基本的な問題で、次のようなコードで解くことができます。(コードはRubyのものです)
dpは「0以上N以下の整数で、条件を満たす場合の数」として定義されます。

N = gets.chomp # Nは文字列で受け取る
K = gets.to_i

# dp[less][k]: ある桁まで見て、N以下で条件を満たすものの
# * less: Nより小さいかフラグ
# * k: 各桁の和
dp = Array.new(2) { Array.new(K+2, 0) }
dp[0][0] = 1

N.size.times do |i|
  ni = N[i].to_i
  dp2 = Array.new(2) { Array.new(K+2, 0) }

  [0, 1].each do |less|
    0.upto(K) do |k|
      0.upto(9) do |d|
        next if less == 0 && d > ni # Nより大きいのでスキップ
        less2 = less
        less2 = 1 if d < ni # Nより小さいことが確定

        k2 = k + d
        k2 = K + 1 if k2 > K # 和がKを超えるものは無視するのでまとめる
        dp2[less2][k2] += dp[less][k]
      end
    end
  end

  dp = dp2
end

ans = 0
[0, 1].each do |less|
  0.upto(K) do |k|
    ans += dp[less][k]
  end
end

# 0の場合を数えているので引く
puts ans - 1

では次の問題です。

問題2

1以上N以下の整数で、各桁の数字の積がK以下のものはいくつあるか?

「各桁の数字の和」が「各桁の数字の積」に変わっただけです。
なので、桁の和を取るところを積に変えて以下のような差分でOKか……というと、これでは正解になりません。

@@ -8 +8 @@
-dp[0][0] = 1
+dp[0][1] = 1 # 積の場合は初期値1から考える
@@ -21 +21 @@
-        k2 = k + d
+        k2 = k * d

なぜでしょうか?

問題1の解答コードでは、Nの各桁ごとに0から9までの数字を当てはめて数を作っています。
つまり、Nが5桁の数であるとき、543という数は543という数字の並びではなく00543という数字の並びで表現されているのです。
和の場合は上の桁に余分な0がついていても影響がありませんが、積の場合はこの余分な0を除外して積を取る必要があります。
この「先頭についている余分な0」のことを英語で「leading zero」と呼びます。

leading zeroを考慮した正しいコードは以下のようになります。
「leading zeroが継続しているか」というフラグを持ち、

  • 数字0が現れたとき、leading zero継続中なら積には加えない
  • 0でない数字が現れたらleading zeroは終了する
    という処理を行います。
N = gets.chomp
K = gets.to_i

# dp[leading_zero][less][k]
# * leading_zero: leading zero中かフラグ
# * less: Nより小さいかフラグ
# * k: 各桁の積
dp = Array.new(2) { Array.new(2) { Array.new(K+2, 0) } }
dp[1][0][1] = 1

N.size.times do |i|
  ni = N[i].to_i
  dp2 = Array.new(2) { Array.new(2) { Array.new(K+2, 0) } }

  [0, 1].each do |leading_zero|
    [0, 1].each do |less|
      0.upto(K+1) do |k|
        0.upto(9) do |d|
          next if less == 0 && d > ni # Nより大きいのでスキップ
          less2 = less
          less2 = 1 if d < ni # Nより小さいことが確定

          leading_zero2 = leading_zero
          leading_zero2 = 0 if d != 0 # 0でない数字が現れたらleading zeroは終わり

          # leading zeroが終わっていたら
          if leading_zero2 == 0
            k2 = k * d
            k2 = K+1 if k2 > K # 積がKを超えるものはまとめてよい
            dp2[leading_zero2][less2][k2] += dp[leading_zero][less][k]
          end
        end
      end
    end
  end

  # 0...0 のときは積は1のものが1通りあるとみなす
  dp2[1][1][1] = 1

  dp = dp2
end

ans = 0
[0, 1].each do |less|
  0.upto(K) do |k|
    ans += dp[0][less][k]
  end
end

# leading_zero == 0のところしか和を取っていないので0の場合は数えていない。
# よってこのままの数字が答えとなる。
p ans

別解法

桁DPの解き方として、「Nより小さいフラグ」を明示的に持たない解法もあります。
この解法では、上の解法と異なりdpを「1以上N未満の整数で、条件を満たす場合の数」として定義します。

この場合、leading zeroフラグを明示的に持つことなくleading zeroが除外されます。

問題によってはこちらの解法の方が解きやすい場合もあります。

N = gets.chomp
K = gets.to_i

# dp[k]: N「未満」で条件を満たす場合の数
# * k: 各桁の積
dp = Array.new(K+2, 0)

# Nと等しい場合の各桁の積
eq_k = 1

N.size.times do |i|
  ni = N[i].to_i
  dp2 = Array.new(K+2, 0)

  # Nより小さい → Nより小さい の遷移
  # Nより小さいことが確定しているので、次はどの数字でも可能
  0.upto(K+1) do |k|
    0.upto(9) do |d|
      k2 = k * d
      k2 = K+1 if k2 > K
      dp2[k2] += dp[k]
    end
  end

  # Nと等しい → Nより小さい の遷移
  # Nの当該桁より小さい数字のみ可能
  0.upto(ni - 1) do |d|
    next if i == 0 && d == 0 # 最初の桁は0にならない
    k2 = eq_k * d
    k2 = K+1 if k2 > K
    dp2[k2] += 1
  end

  # Nと等しい → Nと等しい の遷移
  eq_k *= ni
  eq_k = K+1 if eq_k > K

  # leading zero終了
  if i > 0
    # leading zero終了なので、0..9でなく1..9でループ
    1.upto(9) do |d|
      k2 = d
      k2 = K+1 if k2 > K
      dp2[k2] += 1
    end
  end

  dp = dp2
end

ans = 0.upto(K).sum { |k| dp[k] }

# dpは1以上N未満の数に対して数えているので、
# Nが条件を満たすかどうかは別途数える。
ans += 1 if eq_k <= K

p ans

問題例

  • ABC208 E - Digit Products: 本記事で取り上げたのと同じ問題です。ただしKの上限が大きいため、少し工夫が必要です。公式解説は上の「別解法」の方針で実装しています。
  • ABC194 F - Digits Paradise in Hexadecimal: ある桁までに使用した数字を全て管理しようとすると計算量が厳しいです。lessフラグを明示的に持つより、別解法の方針の方が解きやすいと思います。

Discussion