Maximum sum subset of an array with an extra condition

Question

We are given numbers $n \leq 200$, $k \leq 10$ and an array of $3n$ positive integers not greater than $10^6$. Find the maximum possible sum of a subset of elements of this array, such that in every contiguous $n$ elements there are at most $k$ chosen.

As this is an old high-school level contest problem, I ask for some hints. Also, this means I know there exists a solution far quicker than the proposed by D.W. and most likely not very complicated, so the question is still open.

My ideas mostly involved dynamic programming. I was trying to calculate, for each prefix of the array, best score we can acquire. However, in order to do this, I think I would need to calculate it for every prefix, and every possible choosing of $k$ numbers in the last $n$ numbers of this prefix, resulting in complexity $O(n^{k+1})$, which is far from acceptable. Also, I thought of looking at pairs of positions in the array which are distant by $n$ and their relation to each other, but this approach fails, as in an optimal solution not in every contiguous $n$ element we would choose exactly $k$, sometimes it could be less.

Answer 1

Let's call the length of the array $t \cdot n$, so in your task $t=3$. Since in every fragment $[1, n], [n+1, 2n] \ldots [(t-1)n+1, tn]$ we can't take more than $k$ elements, so the number of taken elements won't be larger than $t\cdot k$. Let's call the indices of taken elements in the given array $x_1, x_2 \ldots x_{tn}$. It's easy to notice that $x_i + n \le x_{i+k}$, because there can't be more than $k$ elements in any window of size $n$.

With this knowledge we are ready to construct dynamic programming solution, which would choose in parallel indices for every segment $[x_1 \ldots x_k], [x_{k+1} \ldots x_{2k}] \ldots [x_{(t-1)k+1} \ldots x_{tk}]$ separately.

$$dp[k+1][k+1][k+1]...(t\,times)[tn]$$ We would iterate trough positions $i$ from $1$ to $tn$ and every time decide whether we want to take one more element in the first segment from the position $i$, in the second segment from the position $i+n$, in the third segment from the position $i+2n$ etc. Bear in mind that at any point any segment cannot have more chosen elements than one of its preceding segments.

For $t = 3$: $$dp[a][b][c][i] = max(dp[a][b][c][i-1], dp[a-1][b][c][i-1] + A[i], dp[a][b-1][c][i-1] + A[n + i], dp[a][b][c-1][i-1] + A[2n + i], dp[a-1][b-1][c][i-1] + A[i] + A[n+i] \ldots)$$ There are in total $3^t$ different combinations of taking elements from $i$-th position and it can be easily fixed in code using additional $3^t$ states. if $a \le b \le c$ is not met: $$dp[a][b][c][i] = -inf$$

Result of the given task is stored in one of $dp[x][0][0][tn], dp[k][x][0], dp[k][k][x]$ states for any $x$ depending on how many elements does the optimal solution use.

So overall time complexity is $O(tn(k+1)^t3^t)$ and space complexity is $O(tn(k+1)^t)$.

Answer 2

Method 1: Merging lazy streams

One approach you could try involves iterating through partial solutions for a part of the array, in decreasing merit, and then merging them:

1. Enumerate the ways of selecting at most $k$ out of the first $n$ elements, in decreasing order of their sum. But don't try to build a list that lists all of these possibilities: instead, build a lazy stream that lets you iterate through this list one by one.

2. Similarly, construct a lazy stream that iterates through the ways of selecting at most $k$ out of the second $n$ elements, in decreasing order of sum; and a lazy stream that iterates through the ways of selecting at most $k$ out of the last $n$ elements.

3. Merge these streams. Basically, given streams $S_1,S_2,S_3$, construct a stream $S$ of elements $s_1+s_2+s_3$, sorted by decreasing order, where $s_1 \in S_1,s_2 \in S_2, s_3 \in S_3$.

4. Iterate through the stream $S$, in decreasing order. Each one gives you a candidate solution. Check whether it satisfies the requirements: i.e., whether for every sequence of $n$ consecutive elements, at most $k$ of them have been chosen. As soon as you have found a single one that satisfies the requirement, you are done; you know that one must be optimal.

So, the problem boils down to how to iterate lazily through these streams. There, I think you can work out a clean way to do that, using priority queues.

There are no guarantees on the running time of this procedure. A worst-case input could force this to run a very long time.

But I expect most inputs will behave better. That's because for many input arrays, if at most $k$ out of the first $n$ elements are selected, and same for the middle $n$ and the last $n$, then there is a fair chance that in every $n$ contiguous elements, at most $k$ are selected.

Method 2: Divide-and-conquer

We can improve method 1 by noting that once we know which ones have been selected from the middle, we can separately find the best for the left side and the best for the right side:

1. Enumerate all ways of choosing at most $k$ out of the middle $n$ elements ($A[n..2n-1]$), ordered by decreasing sum.

2. For each, find the best way to select a subset of the first $n$ elements ($A[0..n-1]$) that can be combined with these middle ones. Separately, find the best way to select a subset of the last $n$ elements ($A[2n..3n-1]$) that can be combined with these middle ones. (These two searches can be done independently.) This gives you a candidate subset of the entire array; compute its sum, and if it is the best seen so far, remember it.

Basically, what we are exploiting is that, given information about which subset of the middle $n$ have been chosen, the subset of the first third is conditionally independent of the subset of the last third. The choice of which ones you select from the first third does not affect your options about which ones you select from the last third.

So, how do we do step 2 efficiently? In other words, given a subset of the middle third of the array, how do we find all ways to extend that to a subset of the first two-thirds that doesn't violate any constraints?

One way is to lazily enumerate all ways of selecting at most $k$ out of the first $n$ elements, in decreasing sum, and then for each test whether it can be combined with the selection you've already got for the middle $n$. This is already enough to do significantly better than method 1. If we have to look $d$ deep into each of the three streams, method 1 will look at $O(d^3)$ combinations, whereas method 2 will look at $O(d^2)$ combinations.

I think the process of selecting the subset of the first third can be sped up a bit further, using your knowledge of the subset of the middle third. Let $i_1,i_2,\dots,i_k$ be the indices of the subset of the middle third that was chosen, where $n \le i_1 < i_2 < \dots < i_k \le 2n-1$. Then this breaks up $A[0..n-1]$ into $k+1$ segments: $A[0..i_1-n]$, $A[i_1+1-n..i_2-n]$, ..., $A[i_{k-1}+1-n..i_k-n]$, $A[i_k+1-n..n-1]$. Let $m_0,m_1,\dots,m_k$ denote the number of items we've selected in each of those segments. Then these numbers must obey all of the following constraints: $m_k = 10-k$, $m_k+m_{k-1} \le 11-k$, $m_k+m_{k-1}+m_{k-2} \le 12-k$, \dots, $m_k+\dots+m_1+m_0 \le 10$. So, enumerate through all of the combinations $m_0,m_1,\dots,m_k$ that satisfy all of these constraints. Then, choose the $m_i$ largest numbers in the $i$th segment of $A[0..n-1]$, and add them to your selection (no need to enumerate over combinations here).