[AtCoder] ABC154 E – Almost Everywhere Zero (500点)

2020年2月10日AtCoder動的計画法, 数え上げ, DP, 500点, 桁DP, 場合分け

問題へのリンク

問題概要

1 以上 N 以下の整数で、0 でない数字がちょうど K 個あるものの個数を求めよ。

制約

\begin{align}
&1 \leq N \leq 10^{100} \\
&1 \leq K \leq 3
\end{align}

考え方

入力の N が非常に大きい問題なので、32bitや64bitの整数型では収まりません。このような場合は、文字列として入力を受け取り、「桁ごとに考える」のが良いです。

また、「1 以上 N 以下の整数について、条件を満たす数」を求めるような問題は、桁DPでうまく計算できることが多いです。

K の値に応じて場合分けをして、ゴリ押しで計算をする方法もありますが、計算が面倒なのでオススメできません。

今回は以下のようなDPを考えましょう。

dp[ i ][ smaller ][ k ] := i 桁目以降で 0 以外の数を使用できるのが残り k 個である数の個数。i 桁目までの部分について、 smaller が true なら N より小さく、false なら N と等しい数であるとする。

解き方

以下のように dp を更新しましょう。

dp[ i ][true][ k ]について

  • dp[i+1][true][ k ] (次の桁が0の時)
  • dp[i+1][true][ k+1 ] (次の桁が0以外の時)

dp[ i ][false][ k ] について

  • dp[i+1][true][ k ] (次の桁が0の時)
  • dp[i+1][true][k-1]*(N[i] – '1’) (次の桁が0でもN[i]でも無い時)
  • dp[i+1][true][k-1] (次の桁がN[i]の時)

実装例

dp が3次元配列となるので、実際にループを回して計算しようとすると頭を使う必要があります。

メモ化再帰で計算することにしました。

DPの別解

以下のようなDPを考えることもできます。

dp[ i ][ smaller ][ k ] := i 桁目までで 0 以外の数を使用したのが k 個である数の個数。i 桁目までの部分について、 smaller が true なら N より小さく、false なら N と等しい数であるとする。

このように考えると、DPの遷移は以下のような場合3つについて考えれば良いです。(参考:桁DP)

  • dp[ i ][true] からはdp[ i+1 ][ture]にのみ遷移( i桁目まででNより小さいなら i+1桁目をどのように選んでもNより小さい)
  • dp[ i ][false] からdp[ i+1 ][ture]へ遷移( i桁目までNと同じで、 i+1桁目はNより小さい数の時)
  • dp[ i ][false] からdp[ i+1 ][false]へ遷移( i桁目までNと同じで、 i+1桁目もNと同じ数の時)

おまけ(場合分けでゴリ押し)

例えば N=314159 とした時には

  • 1~99999
  • 100000~299999
  • 300000~314159

のように場合分けして、さらに K の値で場合分けをすることで、それぞれの場合で条件を満たす数を直接求めることが可能になります。

コーナーケースを色々と考える必要がありかなり面倒ですが、高速に求めることができるのが利点です。

先に、「N に出現する0以外の数」を上位の桁から 3 つ求めておくと、後で計算が少し楽になります。