kmjp's blog

競技プログラミング参加記です

AtCoder ARC #043 : D - 引っ越し

ARC

これは本番中に解きたかったが、Cで時間食いすぎてダメだった。
http://arc043.contest.atcoder.jp/tasks/arc043_d

問題

1～NのN個の家が等間隔で一列に並んでいる。隣接する家同士の距離は1である。
これらの家にM個の家族が住む。各家族はP[i]人構成である。

各家族は異なる家に住まなければならない。
この条件のもとで、これらのM家族における２人対の距離の総和を最大化せよ。

解法

問題設定をN個頂点を距離1の辺でつないだグラフと考える。
すると以前書いたテク、「２点間の距離の総和を考える際は、各辺において両側にある頂点数の積の回数分その辺の長さがカウントされる事を用いる」が利用できるようになる。
TopCoder SRM 662 Div1 Medium ExactTree - kmjp's blog

厳密な証明は公式Editorialに任せるとして、人数の大きな家族をできるだけ両端に置いた方が総和が大きくなることは想像に難くない。
よって、人数の大きな家族の順に、空き家のうち左端と右端どちらに置くか、DPで最大値を求めていく。

上記テクによると、各辺に対し、その辺の左側にx人が住んでいるなら、その辺の分はx*(sum(P)-x)分距離の総和にカウントされる。
言い換えると、各辺で左右に来る人数が確定すればその辺が最終的な距離総和に与える影響は確定できる。

dp[i][x] := i番目までの家族のうち合計x人までが左端寄りの家に配置し、残りを右端寄りの家に配置したときの確定済みの距離の総和
とする。たとえば左端にp家族、右端にq家族いるなら、左端(p-1)辺と右端(q-1)辺の影響が確定する。

i番目の家族までの人数の合計をS(i)とする。

  • i番目の家族を左端に置く場合、左側の辺が新たに１本確定してdp[i][x+P[i]] = max(dp[i][x+P[i]], dp[i-1][x-P[i]] + x*(S(N)-x)で状態を更新する。
  • i番目の家族を右端に置く場合、右側の辺が新たに１本確定してdp[i][x] = max(dp[i][x], dp[i-1][x] + (S(N)-(S(i)-x))*(S(i)-x))で状態を更新する。

M家族配置したら、残された(N-M+1)本の辺の影響はx*(S(N)-x)*(N-M+1)で確定する。

int N,M;
int P[10101];
ll dp[2][100020];
ll sum,tsum;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	FOR(i,M) cin>>P[i], sum+=P[i];
	sort(P,P+M);
	reverse(P,P+M);
	
	MINUS(dp);
	dp[0][P[0]]=0;
	tsum=P[0];
	for(i=1;i<M;i++) {
		int cur=(i-1)%2, tar=cur^1;
		MINUS(dp[tar]);
		FOR(x,100001) if(dp[cur][x]>=0) {
			// left
			dp[tar][x+P[i]]=max(dp[tar][x+P[i]],dp[cur][x]+x*(sum-x));
			// right
			dp[tar][x]=max(dp[tar][x],dp[cur][x]+(sum-(tsum-x))*(tsum-x));
		}
		tsum+=P[i];
	}
	
	ll ma=0;
	FOR(x,100001) if(dp[(M+1)%2][x]>=0) ma=max(ma,dp[(M+1)%2][x]+(N-M+1)*(sum-x)*x);
	cout<<ma<<endl;
}

まとめ

Dとしては難易度がそこまで高いわけはないので、Cが易しければもっと解いた人多いんじゃないかな。