CSES - Coding Company

First, sort the people. This allows us to express the contribution of each team as $(\text{Skill of last person}) - (\text{Skill of first person})$. Let $s_i$ denote the skill of the $i$-th person.

$dp[i][j][k] =$ The number of ways we can form teams from the first $i$ people such that there are $j$ "unfinished" teams and the total penalty so far is $k$. If the $l$-th person is the least skilled in an unfinished team, we define its contribution to the penalty as $-s_l$. Similarly, finishing a team with person $r$ contributes $s_r$ to the penalty. These definitions follow directly from the previous observation that the contribution of a team is $s_r - s_l$.

There are 4 cases when we transition from $i - 1$ to $i$:

• We make a team consisting of only person $i$.

  • $dp[i - 1][j][k]$ transitions to the state $dp[i][j][k]$ (unfinished teams and penalty are not affected).
• We append person $i$ to an unfinished team.

  • This transitions to the state $dp[i][j][k]$ (again, unfinished teams and penalty are not affected). There are $j$ unfinished teams we can choose to extend, so $dp[i][j][k] \mathrel{+}= j \cdot dp[i - 1][j][k]$.
• We use person $i$ to "finish" an unfinished team.

  • This transitions to the state $dp[i][j - 1][k + s_i]$ (one less unfinished team, and we add to penalty as discussed above). There are again $j$ unfinished teams to choose from, so $dp[i][j - 1][k + s[i]] \mathrel{+}= j \cdot dp[i - 1][j][k]$.
• We start a new unfinished team with person $i$.

  • This transitions to the state $dp[i][j + 1][k - s[i]]$ (one more unfinished team, and we subtract from penalty as discussed above).

Two more things:

• $k$ in our DP array can be negative, so just add 5000 to each $k$; there can be at most $N/2 = 50$ unfinished intervals, and each of them contributes at most $-100$.
• $dp[i]$ depends only on $dp[i - 1]$, so we can drop the first dimension that we store (to save memory).

I believe that this is called the "open and close interval trick". See zscoder's blog post for more information and problems.

Copy
#include <bits/stdc++.h>

using namespace std;
using ll = long long;

const int M = 1e9 + 7;
const int K = 5e3;  // K is an offset to account for negatives

int main() {
	int n;