Sum AND Subarrays (AtCoder Dwango Programming Contest V / 第5回ドワンゴからの挑戦状予選 - D)

散々なことをやりまくって時間を溶かしてしまいました

概要

長さNの数列の部分列からK個を選び、それぞれの部分列の和のANDを取ったものの最大値を求める問題。
B - Sum AND Subarrays

解説

ANDをとるので、選んだK個全てで立っているBITのみが残る

ANDの最大値をとるので、最高ビットを残すことが最優先

すべての部分列を列挙( $O(N^2)$ )
上のビットから順番に見ていって、そのビットが立っているような部分列がK個以上あればそのビットを残すことに決める
そして、そのビットが立っている部分列のみをK個の候補に残し、下位のビットに移る

一種の貪欲といえるのかな

部分和の最大値は $10^{12} \simeq 2^{40}$ なので、longを使います
計算量は、部分列が $O(N^2)$ 個、見るビットの数が $O(log(aN))$ 個、僕はsetを使ったのでここに $O(log(N))$ が加わって、 $O(N^2log(aN)log(N))$ でした。
実装を上手くやると $O(N^2log(aN))$ までおちると思います。2桁msで通してる人のコードを読むと大体わかります

ソースコード

▶ソースコードを開く

long n, m;
long a[1005];
long sum[1005];//累積和
long psum[2000000];//部分和
long bitflag = 1;


int main(){
	scanf("%ld%ld", &n, &m);
	for(int i = 1; i <= n; i++){
		scanf("%ld", &a[i]);
		sum[i] = sum[i-1] + a[i];
	}
	long pcnt = 0;
	for(int i = 0; i <= n; i++){
		for(int j = i+1; j <= n; j++){
			psum[pcnt] = sum[j] - sum[i];
			pcnt++;
		}
	}
	long upmost = 0;
	for(long bitflag = 1; bitflag < 35184372088832; (bitflag <<= 1)){
		int cnt = 0;
		for(int i = 0; i < pcnt; i++){
			if(bitflag & psum[i]){
				cnt++;
			}
		}
		if(cnt >= m){
			upmost = bitflag;
		}
	}
	if(upmost == 0){
		printf("0\n");
		return 0;
	}
	std::set<int> st;
	for(int i = 0; i < pcnt; i++){
		st.insert(i);
	}
	while(upmost > 0){
		int cnt = 0;
		for(auto i = st.begin(); i != st.end(); i++){
			if(psum[*i] & upmost){
				cnt++;
			}
		}
		std::set<int> st2;
		if(cnt >= m){
			for(auto i = st.begin(); i != st.end(); i++){
				if(psum[*i] & upmost){
					st2.insert(*i);
				}
			}
			st = st2;
			st2.clear();
		}
		upmost >>= 1;
	}
	long ans = psum[*st.begin()];
	for(auto i = st.begin(); i != st.end(); i++){
		ans = (ans & psum[*i]);
	}
	printf("%ld\n", ans);
}