TopCoder SRM 546 Div1 Medium FavouriteDigits

問題

digit1という数字をcount1個以上含み、
digit2という数字をcount2個以上含む、N以上の最小の数を求めよ。

制約条件

0≦digit1, digit2≦9
count1+count2≦15
Nは10^15以下。
答えは必ずlong longに収まる。

方針

n以下でd1をc1個以上含み、d2をc2個以上含むような数をf(n)と置く。
f(n)は桁DPにより簡単に求まる。

f(k) = f(N-1) + 1となる最小のkを見つけるのが問題なので、
これは二分探索をすればいい。

桁DPは次のような感じ。
dp[先頭からの桁数][digit1の個数][digit2の個数][nより小さいか][leadingzeroが終わったか]
これに、次の桁を10通り試して更新していくようなDP.

ソースコード

int m, d[20];
ll dp[20][20][20][2][2];

ll calc(ll n, int d1, int c1, int d2, int c2){
	for(m = 0; n; n /= 10) d[m++] = n % 10;
	reverse(d, d + m);
	
	memset(dp, 0, sizeof(dp));
	dp[0][0][0][0][0] = 1;
	
	rep(i, m) rep(j, c1 + 1) rep(k, c2 + 1) rep(l, 2) rep(g, 2){
		if(!dp[i][j][k][l][g]) continue;
		rep(h, 10){
			int ng = g || h;
			int nl = l || h < d[i];
			int nj = min(c1, j + (ng && h == d1));
			int nk = min(c2, k + (ng && h == d2));
			if(!nl && h > d[i]) continue;
			dp[i + 1][nj][nk][nl][ng] += dp[i][j][k][l][g];
		}
	}
	
	return dp[m][c1][c2][0][1] + dp[m][c1][c2][1][1];
}

class FavouriteDigits {
	public:
	long long findNext(long long N, int digit1, int count1, int digit2, int count2) 
	{
		ll lo = N - 1, hi = (ll)2e18, mid;
		ll res = calc(N - 1, digit1, count1, digit2, count2) + 1;
		
		while(lo + 1 < hi){
			mid = (lo + hi) / 2;
			if(calc(mid, digit1, count1, digit2, count2) >= res) hi = mid;
			else lo = mid;
		}
		return hi;
	}
};