Codeforces - E. K Balanced Teams

 Problem Link

Problem Summary

There are N students that each have a skill rating a_i. You need to group them (not necessarily all) into at least one and at most K non-empty groups, but the range of skill ratings of the students in a group should not exceed 5. Find the number of ways the students may be divided into groups.

Algorithm Used

Knapsack DP, Sorting

Editorial

First, sort the skill ratings in ascending order - this will make find the answer easier later on.

Let's say we have to find the answer to following problem: make j groups and only choosing from the first i students. To find the answer to this subproblem, we can use the following subproblems:

1. Don't include the ith student in the jth group.
2. Include the ith student in the jth group.

If we don't include the ith student in the jth group, then the answer to this subproblem is simply the answer for grouping i-1students in j group.

But if we include the ith student in the jth group, then we find the number of students that be accomodated in this group. This can be found by binary searching on the sorted array of skill ratings and find the first skill rating that is 5 less the the ith student's skill rating. Let's call the index of this student x. All the skill ratings of the students that come before x will be too small to be included in the jth group. So we find the answer for grouping the students with index less that x, which will be i - x - 1, into j-1 groups and add the number of students in the current group, which will be i - x + 1.

 

Then answer for the inital subproblem (making j groups and only choosing from the first i students) will be the maximum of the above two answers.

DP Formula

\(dp[i][j] = max(dp[i-x-1][j-1] + (i - x + 1), dp[i-1][j]\)

Code

#include <bits/stdc++.h>

using namespace std;

int N, K;

const int MAX_N = 5000 + 1;

int dp[MAX_N][MAX_N];

int nums[MAX_N];

int main() {

    cin >> N >> K;

    for (int i = 1; i <= N; i++)

        cin >> nums[i];

    sort(nums, nums + N + 1);

    nums[0] = -1e9;

    for (int j = 1; j <= K; j++) {

        for (int i = 1; i <= N; i++) {

            int a = 0;

            int x = lower_bound(nums + 1, nums + i, nums[i] - 5) - begin(nums);

            a = dp[x-1][j-1] + (i - x + 1);

            int b = dp[i-1][j];

            dp[i][j] = max(a, b);

        }

    }

    cout << dp[N][K] << '\n';

}