TopCoder SRM 603 Div1 Hard SumOfArrays

問題

長さn項の二つの数列a[i], b[i]が与えられる。
bの順序を好きに入れ替えて数列b'[i]を作る。

c[i] = a[i] + b'[i]という数列に、なるべく同じ数が出るようにbを並べ替える。
このとき、同じ数の出現回数の最大値と、その数を求めよ。
出現回数の最大値になる数が複数ある場合は、もっとも大きい数を求めよ。

制約条件

n≦10万
a[i]≦10万
制限時間4秒

方針

editorial見た。

c[i]に数xが出現する回数の最大値をC[x]とする。
C[x]なる数列を効率的に求めたい。
以下A[i]をa[j] = iなるjの個数、B[i]をb[j] = iなるjの個数とする。

xを固定すると、
C[x] = Σ[i + j = x]min(A[i], B[j])
で、これは謎の発想により見方を変えるとi + j = x, A[i]≧k, B[j]≧kかつk>0である
ペア(i, j, k)の個数に等しい。

このペアを、kの値により場合わけして考える。
kが10以上のとき、A[i]≧kとなるAの要素は1万程度しかないのでループして数えられる。
kが9以下のときは、k=1から9までループをまわす。

C[i + j] += (A[i] >= k) * (B[j] >= k)であるので、
これはすなわち畳み込みで、高速フーリエ変換を使えば高速に計算できる。
制限時間が4秒なのでギリギリ間に合う。

ソースコード

Complex A[SZ], B[SZ];
void calc(int k, const vi &a, const vi &b){
	
	rep(i, MX){
		A[i] = a[i] >= k ? 1 : 0;
		B[i] = b[i] >= k ? 1 : 0;
	}
	for(int i = MX; i < SZ; i++) A[i] = B[i] = 0;
	
	fft(SZ, 2*PI/SZ, A);
	fft(SZ, 2*PI/SZ, B);
	rep(i, SZ) A[i] *= B[i];
	
	fft(SZ, -2*PI/SZ, A);
}

class SumOfArrays {
	public:
	string findbestpair(int n, vector <int> Aseed, vector <int> Bseed) {
		
		vi a = gen(Aseed, n), b = gen(Bseed, n);
		vector<ll> ans(2*MX);
		
		//big k
		vector<pi> biga, bigb;
		rep(i, a.size()){
			if(a[i] >= 10) biga.pb(mp(a[i], i));
			if(b[i] >= 10) bigb.pb(mp(b[i], i));
		}
		rep(i, biga.size()) rep(j, bigb.size()){
			ans[biga[i].second + bigb[j].second] += min(biga[i].first, bigb[j].first) - 9;
		}
		
		//small k
		for(int k = 1; k < 10; k++){
			calc(k, a, b);
			rep(i, 2*MX) ans[i] += (int)(A[i].real() / SZ + 0.5);
		}
		
		pair<ll, int> p(-1, -1);
		rep(i, ans.size()) p = max(p, mp(ans[i], i));
		stringstream ss;
		ss << p.first << " " << p.second;
		return ss.str();
	}
	vi gen(const vi &s, int n){
		vi a, res(MX);
		a.pb(s[0]); a.pb(s[1]);
		rep(i, n - 2){
			a.pb(((ll)a[i + 1] * s[2] + (ll)a[i] * s[3] + s[4]) % s[5]);
		}
		rep(i, n) res[a[i]]++;
		return res;
	}
};