TCO2012 Round2B Hard SequenceTransmission

問題

b + a * i（i = 1, 2, .., n）で表される数列を、二進数で表現し、
初項から隙間をあけずに全部並べて書いた文字列をsとする。

この文字列に含まれる、"01"と、"10"の個数を求めよ。

制約条件

1≦a≦40000
1≦b≦10^18
1≦n≦10^12

方針

桁DP.
a + b ≦ x ≦ a * n + b
なるxののうち、aで割った余りがb % aに等しいもの数に含まれる、
"01" "10"の個数が簡単に求められればよい。
これは、簡単な桁DPによりできる。

lim未満の条件を満たす数をf(lim)とすれば、
求める個数はf(a * n + b + 1) - f(a + b)なので、f(lim)を求める。

dp[pos][mod][nz][prev][s]を、
pos桁目まで考えて、今のところの余りがmodで、
nzは0でない数が出たか、直前の数がprevで、sはlim未満であるか
の状態における、01 10の出現個数とする。

これを更新していくDPをすればよい。
ただし、その状態に至る場合の数が何通りかも覚えておく必要があるので、
way[pos][mod][nz][prev][s]をその状態に至る場合の数として、同時に更新する。

このDPでは、末尾の0によりできる01をカウントしていないので、
それは別個に数える。
末尾の0は、最後の項を除き、必ず01を作るので、末尾の0の数から求めればよい。

ソースコード

ll lo, hi, dp[2][40000][2][2][2], way[2][40000][2][2][2], a, b;
char A[100], B[100];

ll calc(char *d, int n){
	memset(dp, 0, sizeof(dp));
	memset(way, 0, sizeof(way));
	way[0][0][0][0][0] = 1;
	int cur = 0, next = 1;
	rep(i, n){
		memset(dp[next], 0, sizeof(dp[next]));
		memset(way[next], 0, sizeof(way[next]));
		
		rep(mod, a) rep(nz, 2) rep(prev, 2) rep(s, 2) if(way[cur][mod][nz][prev][s]){
			
			rep(j, 2){
				int nmod = (mod * 2 + j) % a;
				int nnz = nz || j;
				int ns = s || d[i] > j;
				if(!ns && j > d[i]) continue;
				
				way[next][nmod][nnz][j][ns] += way[cur][mod][nz][prev][s];
				dp[next][nmod][nnz][j][ns] += dp[cur][mod][nz][prev][s];
				if(nz && prev != j) dp[next][nmod][nnz][j][ns] += way[cur][mod][nz][prev][s];
			}
		}
		swap(cur, next);
	}
	return dp[cur][b % a][1][0][1] + dp[cur][b % a][1][1][1];
}

class SequenceTransmission {
	public:
	long long signalChanges(long long a, long long b, long long N) {
		lo = a + b;
		hi = b + a * N + 1;
		::a = a; ::b = b;
		memset(A, 0, sizeof(A));
		
		int n, m;
		for(m = 0; hi; hi /= 2) B[m++] = hi % 2;
		for(n = m; lo; lo /= 2) A[--n] = lo % 2;
		reverse(B, B + m);
		
		ll ans = calc(B, m) - calc(A, m);
		if(a % 2 == 0) ans += (N - 1) * (b % 2 == 0);
		else ans += (N - 1 + ((a + b) % 2 == 0)) / 2;
		return ans;
	}
};