Need help to optimize dynamic programming problem solution

Question

The problem statement:

Given two arrays a and b with sizes n and m respectively. All numbers in these arrays are in the range of 0 to 9 inclusive. Lets create a matrix with size of n x m where values in row i and column j is equal to ai * 10^9 + bj. Find the path from square 1,1 to n,m with the maxium sum. You're allowed to move forward or down.

Input parameters: The first line contains n and m (1 <= n, m <= 100 000)

The second line contains values of array a

The third line contains values of array b

Output Print the maximum sum

Time limit: 1 second

Memory limit: 512MB

Example:

input:

7 4
0 7 1 7 6 7 6
4 1 9 7

output: 55000000068

I tried to solve this problem with dynamic programming, but my solution works in O(n * m) and can't pass time limit:

#include <iostream>
#include <vector>
#include <cmath>

using namespace std;

int main() {
        int n, m; cin >> n >> m;

        vector<uint64_t> a(n), b(m);
        for(int i = 0; i < n; i++) {
            int tmp;
             cin >> tmp;
             a[i] = tmp * 10e8;
        }

        for(int i = 0; i < m; i++) cin >> b[i];

        vector<uint64_t> dp(m);
        dp[0] = a[0] + b[0];

        for (int i = 1; i < m; i++)
            dp[i] = dp[i-1] + a[0] + b[i];

        for (int i = 1; i < n; ++i) {
            for(int j = 0; j < m; ++j) {
                if (j == 0)
                    dp[j] = dp[j] + a[i] + b[j];
                else
                    dp[j] = max(dp[j], dp[j-1]) + a[i] + b[j];
            }
        }

        cout << dp.back() << endl;

        return 0;
}

Answer 1

It's possible to solve this problem without dynamic programming, and with O(n+m) memory and time requirements.

As pointed out by @Botje in the comments, the reward for higher values in a is overwhelmingly large. An optimal path will therefore remain in the leftmost column until it reaches the largest value in a (which is 7 in the above example). If this maximum value appears only once in a, then the best option would be to consume the whole of this row, followed by the last elements of all the following rows until we reach the bottom right corner.

However, if this maximum value appears more than once, we can get a better score by moving right along the first row with the maximum value of a until we reach a column with the maximum value of b, then moving down this column to the last row containing the maximum value of a. We can then consume the rest of this row followed by the last elements of all following rows as before.

Perhaps an illustration will help:

    a = [ 0, 6, 9, 9, 0, 9, 3, 1 ]
    b = [ 1, 3, 2, 8, 4, 8, 1, 6 ]

   Col:  0     1     2     3     4     5     6     7
  Row:
   0    0,1   0,3   0,2   0,8   0,4   0,8   0,1   0,6 
         v
   1    6,1   6,3   6,2   6,8   6,4   6,8   6,1   6,6 
         v 
   2    9,1 > 9,3 > 9,2 > 9,8   9,4   9,8   9,1   9,6 
                           v
   3    9,1   9,3   9,2   9,8   9,4   9,8   9,1   9,6 
                           v
   4    0,1   0,3   0,2   0,8   0,4   0,8   0,1   0,6 
                           v
   5    9,1   9,3   9,2   9,8 > 9,4 > 9,8 > 9,1 > 9,6 
                                                   v
   6    3,1   3,3   3,2   3,8   3,4   3,8   3,1   3,6 
                                                   v
   7    1,1   1,3   1,2   1,8   1,4   1,8   1,1   1,6 

In this example, there are three rows where a = 9, which are rows 2, 3 and 5. To get the maximum score, we need to follow the first of these rows (i.e. row 2) until we reach the column with the maximum value of b (either column 3 or column 5, it makes no difference). Then move down to the last row where a=9 (row 5), step right to the end of this row, and finally down to the bottom right corner.

I've converted the Python code from an earlier version of this answer into C++. In tests with 10⁵ random values in arrays a and b, it produces a result in about 0.3 seconds on my system. The dynamic programming solution above gives identical results, but takes about 4 minutes:

#include <iostream>
#include <vector>

using namespace std;

int main() {
    int n, m;
    cin >> n >> m;

    vector<int64_t> a(n), b(m);
    int tmp, astart, aend, bmax, i, j;
    
    // Read in arrays a[] and b[]. At the same time,
    // find the first and last indices of the maximum
    // value in a[] (astart and aend) and any index
    // corresponding to the maximum value of b[] (bmax)
    
    for (tmp = -1, i = 0; i < n; i++) {
        cin >> a[i];
        if (a[i] >= tmp) {
            aend = i;
            if (a[i] > tmp) {
                astart = i;
                tmp = a[i];
            }
        }
        a[i] *= 1000000000LL;
    }
    for (tmp = -1, j = 0; j < m; j++) {
        cin >> b[j];
        if (b[j] > tmp) {
            tmp = b[j];
            bmax = j;
        }
    }
    
    // Trace through the matrix. First work down as far as
    // astart, then right until column bmax. Then work down
    // as far as row aend, add the remaining elements in this
    // row, and finally add the last element of each remaining
    // rows until we reach the bottom right corner.
    
    i = j = 0;
    int64_t tot = a[i] + b[j];
    while (i < astart) tot += a[++i] + b[j];
    while (j < bmax) tot += a[i] + b[++j];
    while (i < aend) tot += a[++i] + b[j];
    while (j < m-1) tot += a[i] + b[++j];
    while (i < n-1) tot += a[++i] + b[j];
        
    cout << tot << endl;
    return 0;
}