How can I complete my N*M-element double-sum in O(N+M) time?

Question

I've recently come across an algorithms question involving a function which operates on two arrays and is defined like so:

I implemented the function by using two for loops and adding the values to a variable:

int f(int n, int m, int *a, int *b)
{
    int x = 0;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            x += a[i] * b[j] * (i + j + 2); // (... + 2) to compensate for 0-based indexing
        }
    }    

    return x % 998244353; // constant specified by the exercise
}

The caveat, though, is that the question requires you to implement a solution that runs in linear time in the sizes of A and B, and I couldn't figure out how to condense those two for loops into one, even after thinking about it for quite a bit of time.

Does anyone have any suggestions?

Answer 1

Consider the following two simpler sums:

Bsum = Sum_[j=1...M] B[j] 
      = B[1] + B[2] + ... + B[M]

and

Bjsum = Sum_[j=1...M] j * B[j] 
      =  1 * B[1] + 2 * B[2] + 3 * B[3] + ... + M * B[M]

Note that we can rewrite the formula in your question using these two sums, as follows:

F(A, B) = Sum_[i=1...N] {A[i] * (i * Bsum + Bjsum)}

And thus, what you need to do is compute Bsum and Bjsum, which takes O(M) time, then compute this transformed sum in a straightforward fashion, taking O(N) time. The overall complexity will be O(M)+O(N) - linear in the input size, as required.

Let us now actually write the C code for this computation:

int f(int n, int m, int *a, int *b)
{
    int Bsum = 0;
    int Bjsum = 0;
    for (int j = 0; j < m; j++) {
        Bsum += b[j];
        Bjsum += (j+1) * b[j]);
    }

    int fab = 0;

    for (int i = 0; i < n; i++) {
        fab = (fab + a[i] * ((i+1) * Bsum + Bjsum)) % 998244353;
    }
    return fab;
}