Find the pair that gives the biggest sum such that second element comes after first

Viewed 95

Say I have two numpy arrays of equal length:

a = np.array([1, 3, 5, 3, 2, 0])
b = np.array([1, 1, 2, 2, 4, 3])

Out of all of the i, j such that j >= i, how I can I find the i and j such that

a[i] + b[j]

is maximised?

Given my actual use-case, a tie is very unlikely, but if there is one, returning the smallest i and j would be OK.

6 Answers

Here is an O(n) approach:

A = np.maximum.accumulate(a)
B = np.maximum.accumulate(b[::-1])[::-1]

c = (A+B).argmax()
i = a[:c+1].argmax()
j = c + b[c:].argmax()
i,j
# (2, 4)

How it works:

We first introduce an additional constraint i <= c <= j ; and then get rid of it by maximizing over c.

What we gain is that given the c constraint the best i and j are easily found by looking backward and forward from c.

The last piece of the puzzle is a cumulative maximum trick to compute the best sums for all c in O(n).

Approach #1

We are concerned with all of i, j such that j >= i. That essentially means a upper triangular matrix. So, we need an equivalent mask. Then, we get the pairiwise additions with outer-addition. We mask only the upper triangular ones and look for argmax there and trace back the original (row, col) index pair. Here's a vectorized approach to fulfil those steps -

mask = np.tri(len(a),k=-1,dtype=bool)
minval = min(a.min(),b.min()) # assign as 0 if its all positive numbers
ab = np.where(mask, minval, np.add.outer(a,b))
r,c = np.unravel_index(ab.argmax(),ab.shape)

For given sample, we get -

In [60]: (r,c)
Out[60]: (2, 4)

Approach #2

Another way with np.triu_indices to select based on indices -

R,C = np.triu_indices(len(a))
idx = (a[R]+b[C]).argmax()
r,c = R[idx],C[idx]

Using numba with simple for loops enumerate syntax:

import numpy as np
from numba import jit

a = np.array([1, 3, 5, 3, 2, 0])
b = np.array([1, 1, 2, 2, 4, 3])

@jit(nopython=True)
def compute_sum_ij(a,b):   
    max_sum = -1
    ij = (-1,-1)
    for i,ai in enumerate(a):
        for j,bj in enumerate(b[i:],i):
            ij_sum = ai + bj
            if  ij_sum > max_sum:
                max_sum = ij_sum
                ij = (i,j)           
    return max_sum, ij

returns for given example:

compute_sum_ij(a,b)
>>> 9, (2, 4)

If you care about short answers, here is one using lower triangle of addition a[i]+b[j] matrix:

c = a+b[:,None]
c *= np.tri(*c.shape,dtype=np.int64)
j, i = np.unravel_index(c.argmax(), c.shape)

output:

i, j:
2, 4

If you have negative values: The above solution assumes values are non-negative. If not, simply replace upper triangle of c with -inf. In other words, replace 2nd line above with:

import sys
c *= (np.tri(*c.shape,dtype=np.int64) - sys.maxsize * np.tri(*c.shape,k=-1,dtype=np.int64).T)

If your data type is float, use np.inf instead of large integer sys.maxsize

Explanation:

  • c[j,i] = a[i] + b[j].
  • 2nd line makes sure that c[j,i] for j<i is replaced with -inf (or 0 if values are non-negative), therefore those values are irrelevant in search for maximum.
  • Last line searches for maximum value and returns the corresponding indices.

A straightforward solution would be to use two nested for loops. Something like

import numpy as np

a = np.array([1, 3, 5, 3, 2, 0])
b = np.array([1, 1, 2, 2, 4, 3])

max = -np.inf
max_index_a = -1
max_index_b = -1

for i in range(0, len(a)):
    for j in range(i, len(b)):
        val = a[i] + b[j]
        if val > max:
            max = val
            max_index_a = i
            max_index_b = j

print("maximum at i=" + str(max_index_a) + ", j=" + str(max_index_b) + " with value " + str(max))

Short and simple:

max_val=-np.inf
for ind in range(len(a)-1): 
    curr_max = a[:ind+1].max()+b[ind+1:].max()
    if curr_max >max_val:
        i=a[:ind+1].argmax()
        j=b[ind+1:].argmax()+ind+1
        max_val=curr_max
i,j
#2,4
Related