DP 復元非本質＞＜

問題概要

$N$ 人のコンテスタントを 3 人以上を 1 チームとしたいくつかのチームに分割したい。コンテスタント $i$ のスキルは $a_{i}$ である。

チーム分けの良さを、各チームごとの「メンバーのスキルの最大値と最小値の差」の合計値とする。

チーム分けの良さの最大値を求め、具体的な分け方も求めよ。

制約

$3 \le N \le 2 \times 10^{5}$

考えたこと

とりあえずソートしたくなる。ソート順に区間分割していくイメージで問題なさそう。ここまでは、いつもの、という感じ。区間分割を最適化する DP は基本的には、以下のような定義で $O(N^{2})$ になる。

dp[ i ] := 最初の i 個の要素をいくつかの区間に分割したときの最適値

そして多くの場合、これをどうやって高速化しようか...ということになる。

最適解を変形

今回、区間の長さとして 6 以上のものは考えなくて良い。なぜなら、区間の長さが 6 あったら、それを (3, 3) に分けて損することはないからだ。

よって、区間の長さが 3, 4, 5 のどれかに限って良いことから、計算量は $O(N \log{N}))$ に落ちる (最初のソートがボトルネック)。最後の復元が面倒すぎる。

#include <bits/stdc++.h>
#include <sys/time.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

const long long INF = 1LL<<60;
using pll = pair<long long, int>;
int N;
vector<pll> a;

void solve() {
    vector<long long> dp(N+1, INF);
    vector<int> prev(N+1, -1);
    dp[0] = 0;
    for (int i = 0; i < N; ++i) {
        for (int j = 3; j <= 5 && i + j <= N; ++j) {
            long long add = a[i+j-1].first - a[i].first;
            if (chmin(dp[i+j], dp[i] + add)) {
                prev[i+j] = i;
            }
        }
    }
    int cur = N;
    int id = 0;
    vector<int> res(N, -1);
    while (cur) {
        ++id;
        int nex = prev[cur];
        for (int i = nex; i < cur; ++i) res[a[i].second] = id;
        cur = nex;
    }
    cout << dp[N] << " " << id << endl;
    for (int i = 0; i < N; ++i) {
        if (i) cout << " ";
        cout << res[i];
    }
    cout << endl;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0); 
    while (cin >> N) {
        a.resize(N);
        for (int i = 0; i < N; ++i) {
            cin >> a[i].first;
            a[i].second = i;
        }
        sort(a.begin(), a.end());
        solve();
    }
}