Question

假设我有以下两个numpy数组：

a = numpy.array([[1,4,6,2,5],[3,2,7,12,1],[8,5,3,1,4],[6,10,2,4,9]])
b = numpy.array([0, 1, 4])

现在，我想首先在特定行（例如第二行a[1,:]）中搜索最大值的索引，我有另一个数组b，其中包含一些数字，如果最大值的索引在另一个b中作为元素存在，我使用该索引号进行其他计算。如果a行中不存在来自b行的最大值索引，那么我需要寻找第二大数字的索引，并且如果该索引值存在于{ {1}}作为元素，我继续这样做，否则寻找第三大数字的索引，依此类推。我不想对数组b进行排序。

在上面的示例中，我在第二行中有a。最大数的索引是3，但是b中不存在3，那么第二个最大的索引是2，也不存在在b中，然后我寻找b中存在的第三大索引，即0。因此，然后将0分配给新变量。有任何快捷的方法吗？提前致谢。

Answer 1

这是一个可以很好地扩展通用ndarrays的容器-

def maxindex(a, b, fillna=-1):
    sidx = a.argsort(-1)
    m = np.isin(sidx,b)
    idx = m.shape[-1] - m[...,::-1].argmax(-1) - 1
    out = np.take_along_axis(sidx,idx[...,None],axis=-1).squeeze()
    return np.where(m.any(-1), out, fillna)

样品运行-

In [83]: a
Out[83]: 
array([[ 1,  4,  6,  2,  5],
       [ 3,  2,  7, 12,  1],
       [ 8,  5,  3,  1,  4],
       [ 6, 10,  2,  4,  9]])

In [84]: b
Out[84]: array([0, 1, 4])

In [85]: maxindex(a, b) # all rows
Out[85]: array([4, 0, 0, 1])

In [86]: maxindex(a[1], b) # second row
Out[86]: array([0])

3D外壳-

In [105]: a
Out[105]: 
array([[[ 1,  4,  6,  2,  5],
        [ 3,  2,  7, 12,  1],
        [ 8,  5,  3,  1,  4],
        [ 6, 10,  2,  4,  9]],

       [[ 1,  4,  6,  2,  5],
        [ 3,  2,  7, 12,  1],
        [ 8,  5,  3,  1,  4],
        [ 6, 10,  2,  4,  9]]])

In [106]: maxindex(a, b)
Out[106]: 
array([[4, 0, 0, 1],
       [4, 0, 0, 1]])

Answer 2

IIUC，如果您需要保留a的索引：

res = b[a[:, b].argmax(1)]

array([4, 0, 0, 1])

或者：

a[:, np.delete(np.arange(a.shape[1]), b)] = a.min()
res = a.argmax(1)

array([4, 0, 0, 1], dtype=int64)

如果掩码数组的索引足够：

res = a[:, b].argmax(1)

array([2, 0, 0, 1], dtype=int64)

在多维数组pyhton行中找到最大，第二大或第三大的索引

2 个答案: