kmjp's blog

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

CODE FESTIVAL 2018 Final: G - Chicks and Cages

もうちょいNが大きくてもよかったね。
https://beta.atcoder.jp/contests/code-festival-2018-final-open/tasks/code_festival_2018_final_g

問題

N羽のヒヨコがおり、それぞれ大きさはA[i]とする。
これらのヒヨコをM個の鳥かごに分ける。
各ヒヨコは、同じ鳥かご内にいるヒヨコの大きさの総和の分ストレスを受ける。

最適な分け方を行った場合の総ストレスの最小値を求めよ。

解法

あるヒヨコがn匹いる鳥かごにいる状態のとき、解にはA[i]*nだけ寄与する。
よって小さいヒヨコほどヒヨコの多い鳥かご、大きいヒヨコほどヒヨコの少ない鳥かごに入れるべきである。
そのためヒヨコを大きさ昇順に並べると、各ヒヨコの入る鳥かご内のヒヨコ数は降順になる。

そこで、ヒヨコを大きさ順に並べ、順にグループとして鳥かごに入れていけばよい。
f(i,k) := i羽目までのヒヨコをk個の鳥かごに入れた場合のストレスの総和
上記式を順にDPで求め、最後にf(N,M)が解となる。

事前に大きさの累積和を取っておくと、このf(i,k)はO(N^2*M)で埋められる。
ただヒヨコの羽数が降順であることを用いるともう少し計算量を落とせるが、あいにくそこまでしなくてもよいようだ。

int N,M;
int A[2020];
ll from[2020],to[2020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	FOR(i,N) cin>>A[i];
	sort(A,A+N);
	
	FOR(i,N) from[i+1]=1LL<<60;
	while(M--) {
		FOR(i,N+1) to[i]=1LL<<60;
		FOR(x,N) {
			ll tot=0;
			for(y=x;y<N;y++) {
				tot+=A[y];
				to[y+1]=min(to[y+1],from[x]+tot*(y-x+1));
			}
		}
		
		swap(from,to);
	}
	cout<<from[N]<<endl;
}

まとめ

O(N^2*M)が想定ってことはないよね…。