Codeforces #318 Div1 D. Bear and Cavalry

考察パートが難しい割に、数値の制限がザル。
http://codeforces.com/contest/573/problem/D

問題

軍のN人の兵士がおり、個々の強さはW[i]である。
また各兵士は自分で所有する馬を持ち、その強さはH[i]である。
馬に乗った兵士の強さは、兵と馬の強さの積である。

ここで、各兵士が自分が所有する以外の馬にのるケースを考える。
(兵士と馬は1:1で対応付けられる)
その際、この軍の強さは馬に乗った兵士の強さの和である。

ここでQ個のクエリが与えられる。
各クエリでは、2人の兵士の番号が与えられ、互いの所有する馬を交換する。
各クエリ後の兵士と所有する馬の対応について、うまく兵士と馬の対応を行うことで得られる軍の強さの最大値を求めよ。

解法

兵士が乗る馬に関する制限がない場合、兵士と馬を強さで降順ソートして、i番の兵士がi番の馬に乗ると軍の強さが最大化する。

ここで各兵士が乗れない馬がある場合、そううまい馬とのマッチングができない場合もある。
ただ、厳密な証明はEditorialに任せるとして、軍の強さを最大化させる場合以下が成り立つ。

• i番の兵士が乗る可能性があるのは、(i-2)～(i+2)番の馬である。
• 上の条件より、兵士は強い順に1,2,3人のいずれかで区切っていき、同様に馬を強い順で兵士の人数と同じ間隔で区切っていくと、対応する兵士と馬の間で強さを最大化できるような組み合わせが作れる。

大まかには以下のように考えればよい。
i番までの兵士が、すでにi番までの馬と最適なマッチングが出来ていたとする。
この場合、(i+1)番以降の兵士は以下のようにマッチングすることが考えられる。

• (i+1)番の兵士と(i+1)番の馬をマッチングできるならそうする。
• (i+1)番の兵士と(i+1)番の馬がマッチングできない場合、(i+1)番の兵士が(i+2)番の馬、(i+2)番の兵士が(i+1)番の馬に乗ればよい。
• 基本的には上記のとおり1人か2人でペアを作って馬を交換し合えばよいが、最後に兵士及びマッチング不可な馬が1組余ってしまう可能性がある。その場合どこか1か所3人組を作ればよい。3人組の作り方は以下のいずれか。
  • (i+1)番の兵士が(i+2)番の馬、(i+2)番の兵士が(i+3)番の馬、(i+3)番の兵士が(i+1)番の馬に乗る
  • (i+1)番の兵士が(i+3)番の馬、(i+2)番の兵士が(i+1)番の馬、(i+3)番の兵士が(i+2)番の馬に乗る
  • (i+1)番の兵士が(i+3)番の馬、(i+2)番の兵士が(i+2)番の馬、(i+3)番の兵士が(i+1)番の馬に乗る

上記の事実より、i番目までの兵士と馬の最良なマッチング時の軍の強さから、(i+1),(i+2),(i+3)番目までの兵士と馬の最良マッチング時の軍の強さの候補をDPの要領で求められる。
Editorialには上記DPを毎回繰り返すO(QN)解法では間に合わないとあるが、DP内の分岐を極力減らすとなんと間に合ってしまう。

O(QN)で間に合わないなら、クエリ毎で人と馬の対応が変わるのは一部であることを利用し、SegTreeや平方分割でDPをやり直す範囲を限定すれば間に合う。
以下の回答は手抜きO(QN)。
最初2.9sかかったが、地味に分岐を減らして2.2sにした。

int N,Q;
pair<ll,int> WW[50101];
pair<ll,int> HH[50101];
ll W[50101];
ll H[50101];
int WP[50101],HP[50101],pos[50101];

ll dp2[3][50501];
ll dp[50501];

void update(int i) {
	if(i>N) return;
	dp2[0][i]=dp2[1][i]=dp2[2][i]=-1LL<<50;
	if(i) if(WP[i-1]!=HP[i-1]) dp2[0][i]=W[i-1]*H[i-1];
	if(i>1) if(WP[i-1]!=HP[i-2] && WP[i-2]!=HP[i-1]) dp2[1][i]=W[i-1]*H[i-2]+W[i-2]*H[i-1];
	if(i>2) {
		if(WP[i-1]!=HP[i-2] && WP[i-2]!=HP[i-3] && WP[i-3]!=HP[i-1]) dp2[2][i]=max(dp2[2][i],W[i-1]*H[i-2]+W[i-2]*H[i-3]+W[i-3]*H[i-1]);
		if(WP[i-1]!=HP[i-3] && WP[i-2]!=HP[i-1] && WP[i-3]!=HP[i-2]) dp2[2][i]=max(dp2[2][i],W[i-1]*H[i-3]+W[i-2]*H[i-1]+W[i-3]*H[i-2]);
		if(WP[i-1]!=HP[i-3] && WP[i-2]!=HP[i-2] && WP[i-3]!=HP[i-1]) dp2[2][i]=max(dp2[2][i],W[i-1]*H[i-3]+W[i-2]*H[i-2]+W[i-3]*H[i-1]);
	}
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	FOR(i,N) cin>>WW[i].first, WW[i].second=i;
	FOR(i,N) cin>>HH[i].first, HH[i].second=i;
	sort(WW,WW+N);
	sort(HH,HH+N);
	
	FOR(i,N) W[i]=WW[N-1-i].first, WP[i]=WW[N-1-i].second;
	FOR(i,N) H[i]=HH[N-1-i].first, HP[i]=HH[N-1-i].second, pos[HH[N-1-i].second]=i;
	FOR(i,N) update(i+1);
	
	dp[0]=dp[1]=dp[2]=-1LL<<60;
	
	while(Q--) {
		cin>>x>>y;
		x--,y--;
		swap(HP[pos[x]],HP[pos[y]]);
		swap(pos[x],pos[y]);
		FOR(i,3) update(pos[x]+1+i), update(pos[y]+1+i);
		
		for(i=1;i<=N;i++) dp[i+3]=max(max(dp[i+2]+dp2[0][i],dp[i+1]+dp2[1][i]),dp[i]+dp2[2][i]);
		cout<<dp[N+3]<<endl;
	}
}

まとめ

もう少しQ,Nを大きくしておけば良かったのにね。