lay_contest beta round 00004 神による整地 (3)

問題

www.hackerrank.com

解法

まず基本的な解法ですが,

• 全部同じ値になったらそれが最適
• そうでない場合は, 最大のものと最小のものを選択し, それぞれ -1, +1 する

というのが最適です。この貪欲法を素直にやると 10%, map でやって少し工夫すると 30% の点数がもらえます。

また, 操作は 「+1 を N 回やるパート」と「-1 を N 回やるパート」に分けて, 同じような貪欲をやる, と考えても問題ないことがわかります。満点ではこれを使って考えを少し楽にします。

N がかなり大きい場合, 「すべての数が同じになるまで +1 する」という部分と, 「すべての数に一気に 同じような値を加算する」という部分に分けることが出来ます。後半部分は簡単なのでいいとして, 前半部分が少し工夫が必要なところです。

map でやる場合, 30% 解法から「同じ数を増やす/減らすのは同時にやる」というように操作を変えると, 同じ数を増減させる回数はたかだか 10^5-(-10^5) 回であるので, 十分時間内に答えを得られます。

const int MAX = 555*555;
const ll INF = 1ll<<60;
int RC, H[MAX];

ull calc(ll N) {
    map<ll, int> mp;
    for (int i = 0; i < RC; i++) mp[H[i]]++;
    ll n = N;
    while (1) {
        if (!n) break;
        auto it = mp.rbegin();
        if (it->second == RC) break;
        int num = it->second;
        ll tmp = it->first;
        if (n >= num) {
            n -= num;
            mp[tmp-1] += it->second;
            mp.erase(tmp);
            continue;
        }
        n--;
        mp[tmp-1]++;
        if (--mp[tmp] == 0) mp.erase(tmp);
    }
    if (n) {
        ll p = n/RC;
        int q = n%RC;
        ll tmp = mp.rbegin()->first;
        mp.erase(tmp);
        if (RC-q) mp[tmp-p] = RC-q;
        if (q) mp[tmp-p-1] = q;
    }
    n = N;
    while (1) {
        if (!n) break;
        auto it = mp.begin();
        if (it->second == RC) break;
        int num = it->second;
        ll tmp = it->first;
        if (n >= num) {
            n -= num;
            mp[tmp+1] += it->second;
            mp.erase(tmp);
            continue;
        }
        n--;
        mp[tmp+1]++;
        if (--mp[tmp] == 0) mp.erase(tmp);
    }
    if (n) {
        ll p = n/RC;
        int q = n%RC;
        ll tmp = mp.rbegin()->first;
        mp.erase(tmp);
        if (RC-q) mp[tmp+p] = RC-q;
        if (q) mp[tmp+p+1] = q;
    }
    ull ans = 0;
    for (auto p : mp) {
        ans += (ll)p.first*p.first * p.second;
    }
    return ans;
}

int main() {
    int T;
    cin >> T;
    while (T--) {
        int R, C;
        ll N;
        cin >> R >> C >> N;
        RC = R*C;
        for (int i = 0; i < RC; i++) scanf("%d", H+i);
        cout << min(calc(N), calc(N-1)) << endl;
    }
    return 0;
}