Travelling Salesman and Special Numbers [Codeforces Round #458 C]

http://codeforces.com/contest/914/problem/C

あるxについて以下の変換を考える。
「g(x) := xを2ビット表記したときの1の個数の総和」
２進数表記された数Nが与えられる。
N以下の数でg(x)による変換をK回行うことで丁度1に出来る数は何個あるか(mod10^9+7)

前提知識

• 桁DP

解法

http://codeforces.com/contest/914/submission/34415059

まず気付くべき性質がある
「2^1000以下の数はg(x)を1回通すと1000以下の数になる」
よって、N以下の数について変換を1回通した形、つまり、２進数にしたときの1の数だけを考えればいいことになる。
ここから、以下の桁DPを使う。
「dp[i][sm][isless] := 上位i桁が確定していて、1の数がcnt、未満かどうかがisless、のときの数の個数」
桁DPに慣れていないと難しいが、

        if (N[i] == '0') {
            dp[i + 1][cnt][isless] += dp[i][cnt][isless]; // 末尾に0を置く。islessは変化しない
            if (isless) dp[i + 1][cnt + 1][isless] += dp[i][cnt][isless]; // 末尾に1を置く。これはここまでで未満になってないと置けない。
        } else {
            dp[i + 1][cnt][1] += dp[i][cnt][isless]; // 末尾に0を置く。必ず未満になる。
            dp[i + 1][cnt + 1][isless] += dp[i][cnt][isless]; // 末尾に1を置く。islessは変化しない
        }

この部分が遷移の大事な部分。
「f(x) := g(x)の変換によってxを1に変換する為の変換回数」とすると、全てのcntについてf(cnt)+1==Kとなる数の個数の総和を答える。
　
コーナーケースが結構ある。
K=0のときは1だけなので、答えは1
K=1のときは答えから1を引く必要がある。
これは1のcntは1であるが、これは操作を行っていないのでcntに1が多く含まれることになり、これをカウントしてしまうからである。

string N;
int K;
//---------------------------------------------------------------------------------------------------
int memo[1010];
int f(int x) {
    assert(0 < x);
    if (x == 1) return 0;
    if (memo[x]) return memo[x];

    int y = x, cnt = 0;
    while (y) {
        if (y % 2 == 1) cnt++;
        y /= 2;
    }
    return memo[x] = f(cnt) + 1;
}
//---------------------------------------------------------------------------------------------------
// dp[i][cnt][isless]
// 上位i桁が確定していて、1の数がcnt、未満かどうかがisless、のときの数の個数
mint dp[1010][1010][2];
void _main() {
    cin >> N >> K;

    if (K == 0) {
        cout << 1 << endl;
        return;
    }

    int n = N.length();

    dp[0][0][0] = 1;
    rep(i, 0, n) rep(cnt, 0, n) rep(isless, 0, 2) {
        if (N[i] == '0') {
            dp[i + 1][cnt][isless] += dp[i][cnt][isless]; // put 0
            if (isless) dp[i + 1][cnt + 1][isless] += dp[i][cnt][isless]; // put 1
        } else {
            dp[i + 1][cnt][1] += dp[i][cnt][isless]; // put 0
            dp[i + 1][cnt + 1][isless] += dp[i][cnt][isless]; // put 1
        }
    }

    mint ans = 0;
    rep(cnt, 1, n + 1) rep(isless, 0, 2) if (f(cnt) + 1 == K) ans += dp[n][cnt][isless];
    if (K == 1) ans -= 1;
    cout << ans << endl;
}