Python 找到给出最大和的一对，使第二个元素位于第一个元素之后

假设我有两个长度相等的

numpy

数组：

a=np.array（[1,3,5,3,2,0]）
b=np.数组（[1,1,2,2,4,3]）

在所有

i，j

这样的

j>=i

中，我如何才能找到

i

和

j

a[i]+b[j]

是否最大化

---

考虑到我的实际用例，不太可能有一个tie，但是如果有，返回最小的

i

和

j

就可以了。

一个简单的解决方案是使用两个嵌套for循环。差不多

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))

方法#1

我们关心所有i，j，使得j>=i。这本质上意味着一个上三角矩阵。所以，我们需要一个等效的面具。然后，我们通过

外加法

得到成对加法。我们只屏蔽上面的三角形，在那里查找argmax，并追溯原始（行、列）索引对。下面是实现这些步骤的矢量化方法-

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)

对于给定的样本，我们得到-

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

方法#2

另一种方法是使用

np.triu_index

根据索引进行选择-

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

---

将numba与simple

一起用于

循环

枚举

语法：

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

返回给定示例：

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

以下是一种

O（n）

方法：

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)

工作原理：

---

我们首先引入一个附加约束

i如果您关心简短的答案，这里是一个使用下三角加法
a[i]+b[j]
矩阵的约束：

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

输出：

i, j:
2, 4

如果值为负值：上述解决方案假设值为非负值。如果没有，只需将
c
的上三角替换为-inf。换句话说，将上面的第二行替换为：

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

如果您的数据类型是float，请使用
np.inf
而不是大整数
sys.maxsize

说明：


c[j，i]=a[i]+b[j]

第二行确保
c[j，i]
对于
j简短：

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

如果里面有负数呢？谢谢。容易解决。将在一分钟内添加。认为这对大型阵列有意义。好主意。