AtCoder Grand Contest 018 C - Coins

http://agc018.contest.atcoder.jp/tasks/agc018_c

問題概要

(X+Y+Z)人の人がいて、全員が金・銀・銅の3種類のコインを決められた枚数ずつ持っている。X人から金の、Y人から銀の、Z人から銅のコインをそれぞれその人が持っている分だけもらうとき、もらえるコインの枚数の最大値を求めよ。

解法

まず（金の保有枚数 - 銀の保有枚数）で人を降順にソートする。ソートした列をあるところで前後に区切ると、前の列の人からは金と銅だけ、後ろの列の人からは銀と銅だけもらうことを考えればよくなる（ソート条件より、前の人から銀をもらって後ろの人から金をもらうのは明らかに損なため）。

次に前列から金と銅をどのようにもらうかを考える。前列にK人いるとすると、X人から金のコインをもらい(K-X)人から銅のコインをもらうことになる。ここでK人を改めて（金の保有枚数 - 銅の保有枚数）で降順にソートすると、前半X人から金を、後半(K-X)人から銅をもらうのが最適になる。後列の銀と銅に関しても同様の計算ができる。

最終的な答えは前列と後列の区切り点を全探索し、各区切り点における最適な（金, 銅）, （銀, 銅）の取り方を計算していくことで求めることができる。毎回前列・後列内のソートを行っていては間に合わないので、区切り点をスライドさせながらpriority queue等を使って値を管理していくと計算量がO( (X+Y+Z)log(X+Y+Z) )に抑えられる。

感想

一番最初のソートを入れることで単に2種類の比較に落とせるというのがミソだと思うんですが自力では全く浮かばないですね……

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop;
 
void main() {
    auto s = readln.split.map!(to!int);
    auto X = s[0];
    auto Y = s[1];
    auto Z = s[2];
    auto N = X + Y + Z;
    auto A = new Tuple!(long, long, long)[](N);
    foreach (i; 0..N) {
        s = readln.split.map!(to!int);
        A[i] = tuple(s[0], s[1], s[2]);
    }
    A.sort!"a[1] - a[0] < b[1] - b[0]"();
 
    auto pq = new BinaryHeap!(Array!long, "a < b")();
    auto ans1 = new long[](N);
    long tmp = 0;
 
    foreach (i; 0..N) {
        ans1[i] = tmp;
        tmp += A[i][0];
        pq.insert(A[i][2] - A[i][0]);
        if (i + 1 > X) {
            auto t = pq.front;
            pq.popFront;
            tmp += t;
        }
    }
 
    
    pq.clear;
    auto ans2 = new long[](N);
    tmp = 0;
    
    foreach (i; 0..N) {
        ans2[i] = tmp;
        tmp += A[N-i-1][1];
        pq.insert(A[N-i-1][2] - A[N-i-1][1]);
        if (i + 1 > Y) {
            auto t = pq.front;
            pq.popFront;
            tmp += t;
        }
    }
 
 
    long ans = 0;
    foreach (i; 0..N) {
        if (i >= X && N - i >= Y) ans = max(ans, ans1[i] + ans2[N-i]);
    }
 
    ans.writeln;
}