[CF1256E] Yet Another Division Into Teams

[CF1256E] Yet Another Division Into Teams - dp

Description

有(n)个学生，第(i)个学生的水平为(a_i)。你想把他们分成若干个队伍。每一支队伍至少要有三名学生，总共的极差为所有队伍的极差之和。你需要找到一种最优的分组方案使得总共的极差最小。

Solution

设 (f[i][j]) 表示按权值排序后的前 (i) 个人，最后一组中有 (j) 个人（如果超过 (3) 个当作 (3) 算）

#include <bits/stdc++.h>
using namespace std;

#define int long long

const int N = 1e6 + 5;

int n;
pair<int, int> pr[N];
int f[N][5], p[N][5], a[N], id[N], ans[N];

signed main()
{
    ios::sync_with_stdio(false);
    cin >> n;
    for (int i = 1; i <= n; i++)
    {
        cin >> pr[i].first;
        pr[i].second = i;
    }
    memset(f, 0x3f, sizeof f);
    sort(pr + 1, pr + n + 1);
    for (int i = 1; i <= n; i++)
    {
        a[i] = pr[i].first;
        id[i] = pr[i].second;
    }
    f[1][1] = 0;
    for (int i = 2; i <= n; i++)
    {
        f[i][1] = min(f[i][1], f[i - 1][3]);
        if (f[i][1] == f[i - 1][3])
            p[i][1] = 3;
        f[i][2] = min(f[i][2], f[i - 1][1] + abs(a[i] - a[i - 1]));
        if (f[i][2] == f[i - 1][1] + abs(a[i] - a[i - 1]))
            p[i][2] = 1;
        f[i][3] = min(f[i][3], f[i - 1][2] + abs(a[i] - a[i - 1]));
        if (f[i][3] == f[i - 1][2] + abs(a[i] - a[i - 1]))
            p[i][3] = 2;
        f[i][3] = min(f[i][3], f[i - 1][3] + abs(a[i] - a[i - 1]));
        if (f[i][3] == f[i - 1][3] + abs(a[i] - a[i - 1]))
            p[i][3] = 3;
    }
    int pos = 3, ind = 1;
    for (int i = n; i >= 1; i--)
    {
        ans[id[i]] = ind;
        if (pos == 1 && p[i][pos] == 3)
            ++ind;
        pos = p[i][pos];
    }
    cout << f[n][3] << " " << *max_element(ans + 1, ans + n + 1) << endl;
    for (int i = 1; i <= n; i++)
        cout << ans[i] << " ";
    cout << endl;
}