Given $n$ arrays. Each has size of $h$. Let $a_{i, j} \in \mathbb{I}$ be the $i$-th element of $j$-th array. You can select at most $k$ numbers from all arrays but if you pick $a_{i, j}$, you have to pick $a_{1, j}, a_{2, j}, \dots, a_{i-1, j}$ as well. Each element can only be picked once. Find maximum sum of elements.

My greedy solution is : Let $p_{i, j}$ be $\sum_{x=1}^{i} a_{x, j}$. Sort a list of $\frac{p_{i, j}}{i}$ for all $i$ and $j$ in descending order. Then iterate through the list and check each value whether we have included the prefix sum of this array into our final answer or not. If no and $h \geq j$, we include it in the final answer, subtract $i$ from $h$, and mark $j$ as included. If yes and $j >$ the index of the previous one, we do the same thing just the part extended from the previous one.

Is this solution optimal or is there a better solution to this problem?

  • $\begingroup$ Can you pick the same number multiple times? $\endgroup$
    – orlp
    Feb 2, 2019 at 14:19
  • $\begingroup$ No, you can only pick each number once. $\endgroup$
    – PeppaPig
    Feb 3, 2019 at 1:05
  • $\begingroup$ cs.stackexchange.com/q/59964/755 $\endgroup$
    – D.W.
    Feb 3, 2019 at 2:55

1 Answer 1


Unfortunately, your greedy algorithm may return elements whose sum is not the maximum. For example, let $n=h=k=2$ and $$\begin{pmatrix} a_{11} & a_{12}\\ a_{21} & a_{22} \end{pmatrix} = \begin{pmatrix} 3 & 1\\ 2 & 1 \end{pmatrix}.$$ Your algorithm will pick 3 and 1. The optimal elements should be 3 and 2.

The natural approach to this problem is dynamic programming.

Here are the subproblems.

For nonnegative integers $t$ and $j$ such that $0\le j\le n$ and $0\le t \le \min(nh, t)$, find $dp(t,j)$, the maximal sum of $t$ elements from the first $j$ arrays where all elements selected from the same array must be a prefix of that array. The wanted maximum sum is the maximum number of all $dp(t,j)$.

In order to select $t$ elements from the first $j$ arrays, we can select $i$ elements from the $j$-th array and $t-i$ elements from the first $j-1$ arrays for some $i$, $0\le i\le \min(t,h)$. Hence we have the following recurrence relation.

$$ dp(t,j) = \max_{0\le i\le \min(t,h)}\left(dp(t-i, j-1) + \sum_{1\le k\le i}a_{k,j}\right)$$

The base cases and the boundary cases can be figured out easily.

  • $\begingroup$ A nice solution to a nice problem. I have passed some time on this problem wondering if it was polymonial. And if I am correct, this is a O(kn²). $\endgroup$
    – Optidad
    Feb 4, 2019 at 14:06

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.