読者です 読者をやめる 読者になる 読者になる

torus711 のアレ

主に競技プログラミングの問題について書きます

AtCoder Biginner Contets 029, D - 1

問題概要

 正整数 $N$ が与えられる.$N$ 以下の正整数全てを 10 進数で表したとき,出現する 1 の個数を求めよ.
 $N < 10^9$

解法

 $N$ 以下の整数全てを列挙して 1 の個数を数えようとすると $O( N \log_{ 10 } N )$ 時間かかります.これだと恐らく TLE なのでより効率的な方法を考えます.
 ナイーブに書くと「$N$ 以下の整数を列挙」する処理が必要になる問題の場合,Digit DP が有効な場合がしばしばあります.この問題も,Digit DP で解くことができます.
 次のように状態をとります.

  • dp[ 処理した桁数 ][ 1 の個数 ][ $N$ 未満確定か ] := この状態に至るパターン数

初期状態は dp[ 0 ][ 0 ][ 0 ] のみ 1 で,それ以外は 0 とします.dp[ i ][ j ][ k ] からの遷移は次の桁に入れる数字を全て試します.$N$ に於ける着目している桁の数字を $D$ ,次に入れようとしている数字を $d$ とすると,DP の更新は

  • dp[ i + 1 ][ j + ( d == 1 ) ][ k || d < D ] に dp[ i ][ j ][ k ] を足し込む

となります.$d$ として選べるのは,$N$ 未満確定のとき(⇔ k == 1 )のとき $0$ 以上 $9$ 以下の全ての値で,そうでないときは $0$ 以上 $D$ 以下の数値です.
 この計算が終わった後で,dp[ |S| ][ j ][ k ] で表される状態について j * dp[ |S| ][ j ][ k ] の総和を求めれば,それが問題の答えです.
 (有効な)DP の状態数が $O( \log_{ 10 }^2 N )$ であり,状態遷移は定数個なので,この DP は $O( \log_{ 10 }^2 N )$ 時間で完了します.

コード

using LL = long long;

#define REP2( i, n ) REP3( i, 0, n )
#define REP3( i, m, n ) for ( int i = ( int )( m ); i < ( int )( n ); ++i )
#define GET_REP( a, b, c, F, ... ) F
#define REP( ... ) GET_REP( __VA_ARGS__, REP3, REP2 )( __VA_ARGS__ )

#define SZ( v ) ( (int)( v ).size() )

LL dp[ 16 ][ 16 ][ 2 ];

int main()
{
	cin.tie( 0 );
	ios::sync_with_stdio( false );

	string S;
	cin >> S;

	const int L = SZ( S );
	
	dp[0][0][0] = 1;
	// dp[ digits ][ ones ][ less N ]

	REP( i, L )
	{
		const int D = S[i] - '0';

		REP( j, L )
		{
			REP( k, 2 )
			{
				REP( d, k ? 10 : D + 1 )
				{
					dp[ i + 1 ][ j + ( d == 1 ) ][ k || d < D ] += dp[i][j][k];
				}
			}
		}
	}

	LL res = 0;
	REP( j, L + 1 )
	{
		REP( k, 2 )
		{
			res += dp[L][j][k] * j;
		}
	}
	cout << res << endl;

	return 0;
};