AtCoder ABC #252 : Ex - K-th beautiful Necklace

TLEを解消しきれず…。
https://atcoder.jp/contests/abc252/tasks/abc252_h

問題

C個の数列が与えられる。全数列の要素数の総和をNとする。
各数列から1つずつ要素を選び、そのxorを取った値を列挙した場合、(同じ値の重複もカウントしたとして)大きい順にK番目に来る値は何か。

解法

要素の選び方は、最大で3^(N/3)通り程度になる。これは10^11ちょいである。
そこで、半分全列挙することにしよう。
いくつかの数列から要素を選んだxorを取った値の数列をA、残った数列から要素を選んだxorを取った値の数列をBとする。

A

と

B

はだいたい同じサイズになるようにしておこう。

あとは二分探索で「a∈|A|とb∈|B|を選んだ時a xor bがV以上の値を取るものはK個以上ある」というVを求めよう。
典型的な手法としては、Bに対しTrieを作っておき、a∈|A|に対しa xor bがV以上となるbを数え上げる手段がある。
ただ今回のケースだと微妙にTLに収まりきらない。

そこで、明にTrieを作るのではなく、以下のようにする。
Vを上の桁から決めるよことを考える。

この時、数列のペアの数列Pを以下のようにとる。

• Vのうち確定済みのbitにおいて、a∈|A|とb∈|B|のうち、a xor bがその確定済みbitに一致し、かつaの該当する上位bitの位置が同じものだけからなる数列A'とbの該当する上位bitの位置が同じものだけからなる数列B'のペア(A',B')の配列

この状態は、a'∈|A'|とb'∈|B'|について、a' xor b'が真の解より大きいか小さいか未確定な状況であるものを残していることに相当する。
あとは上から順にVを定めていこう。

int N,C;
vector<ll> V[70];
ll K;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>C>>K;
	FOR(i,N) {
		ll v;
		cin>>x>>v;
		V[x-1].push_back(v);
	}
	
	vector<ll> A,B;
	A={1},B={1};
	vector<pair<int,int>> cand;
	FOR(i,C) cand.push_back({-V[i].size(),i});
	sort(ALL(cand));
	FORR2(a,i,cand) {
		if(A.size()>B.size()) swap(A,B);
		vector<ll> D;
		FORR(a,A) FORR(v,V[i]) D.push_back(a^v);
		swap(A,D);
	}
	
	ll ret=0;
	vector<pair<vector<ll>,vector<ll>>> D={{A,B}};
	
	for(i=59;i>=0;i--) {
		vector<pair<array<vector<ll>,2>,array<vector<ll>,2>>> nex;
		ll num=0;
		FORR2(a,b,D) {
			array<vector<ll>,2> na,nb;
			FORR(v,a) na[(v>>i)&1].push_back(v);
			FORR(v,b) nb[(v>>i)&1].push_back(v);
			nex.push_back(make_pair(na,nb));
			num+=na[0].size()*nb[1].size()+na[1].size()*nb[0].size();
		}
		
		D.clear();
		if(num<K) {
			//このbitが1だと条件を満たせない
			K-=num;
			FORR2(a,b,nex) {
				if(a[0].size()&&b[0].size()) D.push_back({a[0],b[0]});
				if(a[1].size()&&b[1].size()) D.push_back({a[1],b[1]});
			}
		}
		else {
			//このbitが1となるものの中に解がある
			ret+=1LL<<i;
			FORR2(a,b,nex) {
				if(a[0].size()&&b[1].size()) D.push_back({a[0],b[1]});
				if(a[1].size()&&b[0].size()) D.push_back({a[1],b[0]});
			}
		}
		
	}
	cout<<ret<<endl;
	
}

まとめ

Trieを作った方が速いと思ったけど、数列を切り分けていく方が速かったか…。