TopCoder SRM 605 Div1 Medium AlienAndSetDiv1

問題

自然数N, Kが与えられる。
{1, 2, ..., 2*N}の集合を互いに素なN個ずつの集合A, Bにわける。
A, Bの要素を小さい順に並べたとき、どのiについても|A[i] - B[i]| ≧ KとなるようなA, Bの選び方は何通りあるか、mod 10^9 + 7で求めよ。

制約条件

N≦50
K≦10

方針

O(N^2)でできるじゃーんとかいう電波を受信してハマってた。

A, Bに1から順に要素を詰めていくDP.
今のところのAの要素をa個、Bの要素をb個とする。
常にa≧bであるとしてよい。

問題なのはAの新しいほうからa-b個の部分で、この部分をbitで表現できればいい。
集合の選び方への干渉がありうるのは、今追加する数をxとしたら、xとの差が1〜k-1であるもの。
なのでa-b個に入ってるもののうち、差が1〜k-1であるようなものだけbitで持てばいい。

このbitに、1がa-b個立っていたら、Bへの追加はできないということ。

計算量はO(N^2 logK)なんだけど、K-1桁分しか見ないし、a≧bで1/2がかかるし、
実際に遷移する状態もあんまり多くないしで、mapを使っても0.5s弱くらいで余裕で通る。

ソースコード

const int mod = (int)1e9 + 7;
int n, k, mask;
map<pair<pi, int>, int> dp;

int rec(int a, int b, int bit){
	if(a == n && b == n) return 1;
	if(a > n) return 0;
	if(dp.count(mp(mp(a, b), bit))) return dp[mp(mp(a, b), bit)];
	
	int &res = dp[mp(mp(a, b), bit)];
	
	res = rec(a + 1, b, (bit << 1 | 1) & mask);
	if(a == b) res *= 2;
	if(__builtin_popcount(bit) < a - b) res += rec(a, b + 1, (bit << 1) & mask);
	
	return res %= mod;
}

class AlienAndSetDiv1 {
	public:
	int getNumber(int N, int K) {
		n = N; k = K; mask = (1 << k - 1) - 1; //1〜k-1だけ覚える
		dp.clear();
		return rec(0, 0, 0);
	}
};