Codeforces #955 : Div2 E. Number of k-good subarrays

ちょっと手間取ったけどまぁどうにか。
https://codeforces.com/contest/1982/problem/E

問題

整数N,Kが与えられる。
区間[L,R]がK-goodであるとは、L以上R以下の整数が、いずれも二進数表記で1の数がK以下であることをいう。
N以下の非負整数からなる区間で、K-goodであるものは何通りか。

解法

区間[L,R]がK-goodであることは、数列(L,L+1,...,R)がいずれも二進数表記で1の数がK以下であることと同じ。

以下を考える。
f(N,K) := 0～N未満の整数を順に並べた数列において、2進数表記でK以下の数が、prefix/suffixで何個続くか。また、それ以外にこの数列内で二進数表記で1の数がK以下であるものは何通りか。

N以下の最小の2のべき乗の数をMとする。
区間にMとM+1を含まないものは、f(M,K)とf(N-M,K-1)それぞれにおいて、条件を満たす数列の数の和で求められる。
また、MとM+1を含むものは、f(M,K)とf(N-M,K-1)のprefix/suffixで条件を満たす数列が何個あるかから計算できる。

int T;
ll N,K;
const ll mo=1000000007;

map<ll,array<ll,4>> memo[66];

array<ll,4> hoge(ll R,int k) {  // L,R,all,sum
	if(R==0||k<0) return {0,0,0,0};
	if(R<=1LL<<k) return {0,0,R%mo,0};
	if(k==0) return {1,0,0,0};
	
	if(memo[k].count(R)) return memo[k][R];
	
	ll M=1;
	while(M*2<R) M*=2;
	auto a=hoge(M,k);
	auto b=hoge(R-M,k-1);
	array<ll,4> c={0,0,0,0};
	if(a[2]&&b[2]) {
		c[2]=a[2]+b[2];
	}
	else if(a[2]) {
		c[0]=a[2]+b[0];
		c[1]=b[1];
		c[3]=b[3];
	}
	else if(b[2]) {
		c[0]=a[0];
		c[1]=a[1]+b[2];
		c[3]=a[3];
	}
	else {
		c[0]=a[0];
		c[1]=b[1];
		ll v=(a[1]+b[0])%mo;
		c[3]=a[3]+b[3]+v*(v+1)/2;
	}
	
	
	c[0]%=mo;
	c[1]%=mo;
	c[2]%=mo;
	c[3]%=mo;
	memo[k][R]=c;
	return c;
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>K;
		auto a=hoge(N,K);
		a[3]+=a[0]*(a[0]+1)/2;
		a[3]+=a[1]*(a[1]+1)/2;
		a[3]+=a[2]*(a[2]+1)/2;
		cout<<a[3]%mo<<endl;
		
	}
}

まとめ

これは方針立ちやすかったしね。