yukicoder : No.3605 Grand Cross

うまく書かないとコードが長くなりそう。
https://yukicoder.me/problems/no/3605

問題

N要素の整数列AとM要素の整数列Bが与えられる。
また、各要素には色が指定される。

A,Bのうち連続するn要素を抽出したとする。
その際、nが奇数で、かつAとBから抽出した要素の真ん中の位置の色が一致するようにしたい。
抽出した要素の最大値を求めよ。

解法

nは極力大きい方がよい。
よって色が一致する要素を固定したとき、A[0],A[N-1],B[0],B[M-1]のいずれかは抽出した要素に含まれるはず。
そこで以下はA[0]が含まれる場合を求める。(あとはAとBを入れ替えたり、A,Bを反転させながら試せばよい)

A[0]を含むためには、抽出する真ん中の要素はAの前半分となる。
まず真ん中の要素を総当たりし、色ごとにAから抽出できる要素数とその時の和を求めておく。

次にBから抽出した真ん中の要素を総当たりする。
「色ごとにAから抽出できる要素数とその時の和」を二分探索することで、Bの真ん中の要素ごとに、nをどこまで伸ばせるかわかる。

int T;
int N,M;
ll maR[404040],maC[404040];
ll As[202020],Bs[202020];
vector<int> cand[404040];

ll hoge(vector<int> A,vector<int> B,vector<int> C,vector<int> D) {
	int N=A.size(),M=B.size();
	int i,j;
	ll ret=-1;
	FOR(i,N) {
		As[i+1]=As[i]+A[i];
		cand[C[i]].push_back(i);
	}
	FOR(i,M) Bs[i+1]=Bs[i]+B[i];
	FOR(i,M) {
		int p=lower_bound(ALL(cand[D[i]]),min({i,M-1-i,(N-1)/2})+1)-cand[D[i]].begin();
		if(p) {
			p=cand[D[i]][p-1];
			ret=max(ret,As[2*p+1]+Bs[i+p+1]-Bs[i-p]);
		}
	}
	
	FOR(i,N+M) cand[i].clear();
	return ret;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>M;
		vector<int> A(N),B(M),C(N),D(M);
		FOR(i,N) cin>>A[i];
		FOR(i,M) cin>>B[i];
		FOR(i,N) {
			cin>>C[i];
			C[i]--;
		}
		FOR(i,M) {
			cin>>D[i];
			D[i]--;
		}
		ll ret=-1;
		ret=max(ret,hoge(A,B,C,D));
		ret=max(ret,hoge(B,A,D,C));
		reverse(ALL(A));
		reverse(ALL(B));
		reverse(ALL(C));
		reverse(ALL(D));
		ret=max(ret,hoge(A,B,C,D));
		ret=max(ret,hoge(B,A,D,C));
		cout<<ret<<endl;
		
	}
}

まとめ

こういうの迷わずすっきり書けると良いんだけどな。