AtCoder Biginner Contets 029, D - 1

問題文

　http://abc029.contest.atcoder.jp/tasks/abc029_d

問題概要

　正整数 $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;
};