torus711 のアレ

主に競技プログラミングの問題について書きます.PC 以外だと数式が表示されないかもしれないです

AtCoder Beginner Contest 336, E : Digit Sum Divisible

問題概要

 正整数 $n$ が与えられる.$n$ 以下の正整数で,その桁和が自身を割り切るものはいくつあるか?

制約

  • $1 \leq n \leq 10^{ 14 }$

解法

 この問題の制約下では,$n$ 以下の正整数の桁和の最大値は $9$ を目一杯並べたときにできる $9 \times 14 = 216$ です.これは十分に小さいので,桁和ごとに問題を分割し,それぞれの答えの和をとることで元の問題を解くことができそうです.
 では,桁和を固定した場合について考えていきます.固定した桁和を $m $ とします.このとき,条件を満たす正整数は,

  • 桁和が $m $ である.
  • $m $ での剰余が $0$ である.

の両方を満たします.こういった,「ある値以下の非負整数*1の内,いくつかの条件を満たすものを数え上げる」問題を解く手法のひとつとして Digit DP があり,今回の問題も Digit DP で解くことができます.具体的には,先頭から一桁ずつ値を決めていくことを考えて,次のように状態をとる DP をします:$$\mathit{ dp }[i][j][k][l] := \text{# of ways}$$
ここで,添字 $i, j, k, l$ の意味は,

  • $i$ : 決めた桁数.
  • $j$ : 桁和.
  • $k$ : $m $ での剰余.
  • $l$ : $l = 0$ のとき,構成中の数は $n$ に等しい.$l = 1$ なら $n$ 未満.

です.「配る DP」を考えたとき,$\mathit{ dp }[i][j][k][l]$ からの遷移は次の桁に入れる数 $d’$ を全て試すことで,$n$ の上から $i$ 桁目 ($0$-based) を $d$ とすると,$d'$ の範囲は
\begin{align*}
1 \leq d' \leq d && \text{($l = 0$)} \\
1 \leq d' \leq 9 && \text{($l = 1$)}
\end{align*}
となります.遷移先の状態の添字 $i', j', k', l'$ は,
\begin{align*}
i' &= i + 1 & \text{(1 桁増えるので)} \\
j' &= j + d' & \text{(桁和に $d'$ が足されるので)} \\
k' &= ( 10 k + d' ) \bmod m & \text{(構成中の数が一桁シフトしてから最下位桁が $d'$ になるので)} \\
l' &= l \lor ( d' < d )
\end{align*}
と求まります.この DP を回し終わった後,$n$ の桁数を $\mathit{ len }( n )$ として,$$\mathit{ dp }[ \mathit{ len }( n ) ][m][0][0] + \mathit{ dp }[ \mathit{ len }( n ) ][m][0][1]$$ が部分問題の答えです.あとはあり得る $m $ 全てについて和をとれば,元々の問題の答えが得られます.

計算量について

 上記 DP の計算量について考えます.状態数について考えると,($10$ 進数の $10$ は定数だと思えば*2,)$i, j, k$ の値はそれぞれ $\Theta( \log_{ 10 } n )$ 通りで,$l$ は $2$ 通りなので,状態数=空間計算量は $\Theta( ( \log n )^3 )$ 空間となります.そのそれぞれから定数個の状態遷移があるので,時間計算量も $\Theta( ( \log n )^3 )$ 時間です.
 この DP を合計 $\Theta( \log n )$ 回回すので,全体では $\Theta( ( \log n )^4 )$ 時間となります.

コード

LL solve( const string &S, const int M )
{
	const int L = SZ( S );

	static LL dp[ 16 ][ 1 << 7 ][ 1 << 7 ][2];
	// dp[ # of decided digits ][ digit sum ][ modulo M ][ less than N ? ]
	fill( AALL( dp ), 0 );
	dp[0][0][0][0] = 1;

	REP( i, L )
	{
		REP( j, M + 1 )
		{
			REP( k, M )
			{
				REP( l, 2 )
				{
					const int D = S[i] - '0';

					REP( d, l ? 10 : D + 1 )
					{
						if ( M < j + d )
						{
							continue;
						}

						dp[ i + 1 ][ j + d ][ ( k * 10 + d ) % M ][ l || d < D ] += dp[i][j][k][l];
					}
				}
			}
		}
	}

	return dp[L][M][0][0] + dp[L][M][0][1];
}

int main()
{
	IN( string, S );

	LL res = 0;
	REP( i, 1, 1 << 7 )
	{
		res += solve( S, i );
	}

	cout << res << endl;

	return 0;
}

*1:「非負」としてますが,後述の戦略から $0$ は除外されるので,Digit DP の性質に合わせて非負と言っています.

*2:これは許してほしい.