Greedy Solution for Selecting Prefix Sum

Question

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?

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.