AtCoder ARC #127 : D - Sum of Min of Xor

開始時間の変更に気付かず不参加だった回。
https://atcoder.jp/contests/arc127/tasks/arc127_d

問題

2つの整数列A,Bが与えられる。
各添え字の対(i,j)に対し、min(A[i]^A[j],B[i]^B[j])の総和を求めよ。

解法

2進数で上の桁からA[i],B[i],A[j],B[j]の各bitを比較したとき、(A[i]^B[i])と(A[j]^B[j])の値がずれたところで、A[i]^A[j]とB[i]^B[j]のどちらが小さくなるか確定する。
そこで、(A[i]^B[j])の上d桁が一致するような(A[i],B[i])の対をグループにして処理していこう。

int N;
int A[252525],B[252525];
ll ret=0;

void dfs(vector<pair<int,int>> V,int d) {
	if(V.size()<=1) return;
	if(d<0) {
		ll NA[20][2]={};
		FORR2(a,b,V) {
			int i;
			FOR(i,20) {
				int x=(a>>i)&1;
				ret+=NA[i][x^1];
				NA[i][x]+=1<<i;
			}
		}
		return;
	}
	
	vector<pair<int,int>> W[2];
	FORR2(a,b,V) {
		W[((a^b)>>d)&1].push_back({a,b});
	}
	ll NA[2][2][20][2]={};
	ll NB[2][2][20][2]={};
	int i;
	FORR2(a,b,W[0]) {
		int ad=(a>>d)&1;
		int bd=(b>>d)&1;
		FOR(i,19) {
			NA[ad][bd][i][(a>>i)&1]+=1<<i;
			NB[ad][bd][i][(b>>i)&1]+=1<<i;
		}
	}
	
	FORR2(a,b,W[1]) {
		int ad=(a>>d)&1;
		int bd=(b>>d)&1;
		FOR(i,19) {
			ret+=NA[ad][bd^1][i][1^((a>>i)&1)];
			ret+=NB[ad^1][bd][i][1^((b>>i)&1)];
		}
	}
	
	dfs(W[0],d-1);
	dfs(W[1],d-1);
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	FOR(i,N) cin>>B[i];
	vector<pair<int,int>> V;
	FOR(i,N) V.push_back({A[i],B[i]});
	dfs(V,19);
	cout<<ret<<endl;
}

まとめ

割とコード量が増えてしまった。
もっとすっきり書く方法あるのかな。