K Balanced Teams（31-33）

开启掘金成长之旅！这是我参与「掘金日新计划 · 2 月更文挑战」的第 31 天，点击查看活动详情

Problem - E - Codeforces

原文题面

题目描述

有 $n$ 个学生和 $k$ 支队伍，队伍中的每个人具有一个数值属性，每个队伍中任意两个人的属性值相差不超过 $5$。

在最大限制队伍数为 $k$ 的情况下，问 $k$ 支队伍中最多可以包含多少学生。

数据范围

$1≤k≤n≤5000$

$1≤a_i≤10^9$

测试样例

样例1

input

5 2
1 2 15 15 15

output

5

样例2

input

6 1
36 4 1 25 9 16

output

2

样例3

input

4 4
1 10 100 1000

output

4

题目分析

这是一道单调队列优化线性DP的题目。

首先我们考虑暴力做法，每个学生具有两种状态，进入某支队伍或不进入队伍，枚举的复杂度为 $2^n$，这显然是不行的。

考虑动态规划。

首先我们将 $n$ 位同学按属性数值大小排序。

定义 f[i][j] 表示前 $i$ 位同学共分出了 $j$ 支队伍的包含最大人数。

对于第 $i$ 位同学，考虑两种情况，即选或不选。

若不选择， f[i][j] = f[i-1][j]。

若选择，我们假设第 $i$ 位同学的数值为 $s$，并以他为其所在队伍的数值最大值，假设属性数值小于当前队伍的上一个队伍的最后一个人的位置为 $l$，$l$ 位置上的数值 $c$ 为小于 $s-5$ 的某个值。

转移方程为 f[i][j] = f[l][j-1] + i - l + 1。

可以维护一个单调队列表示以 $i$ 为队尾的学生标号，若队头学生的数值小于 $s-5$，则将其弹出。

注意最开始的初始化为 f[1][1~n] = 1，即存在人数为 $0$ 的队伍，这样写是为了贴合代码的转移方式。

Accept代码 $O(n+k)$

#include <bits/stdc++.h>

using namespace std;

const int N = 5010;

int n, k;
int a[N];
int f[N][N];

int stk[N], hh, tt = -1;

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);

    cin >> n >> k;
    for (int i = 1; i <= n; i ++) cin >> a[i];
    sort(a + 1, a + 1 + n);

    for (int i = 1; i <= k; i ++) f[1][i] = 1;
    stk[++ tt] = 1;
    for (int i = 2; i <= n; i ++)
    {
        stk[++ tt] = i;
        while (a[stk[tt]] - a[stk[hh]] > 5) hh ++;
        for (int j = 1; j <= k; j ++)
        {
            f[i][j] = f[i - 1][j];
            f[i][j] = max(f[i][j], f[stk[hh] - 1][j - 1] + i - stk[hh] + 1);
        }
    }
    cout << f[n][k];
    return 0;
}