1

我有一个 2D numpy 数组，我需要沿特定轴取最大值。然后我需要知道为该操作选择了哪些索引作为另一个操作的掩码，该操作仅在那些相同的索引上完成，但在另一个相同形状的数组上完成。

正确我是如何通过使用二维数组索引来做到这一点的，但它很慢而且有点复杂，尤其是生成行索引的 mgrid hack。这个例子只是 [0,1] 但我需要鲁棒性来处理任意形状。

a = np.array([[0,0,5],[0,0,5]])
b = np.array([[1,1,1],[1,1,1]])
columnIndexes = np.argmax(a,axis=1)
rowIndexes = np.mgrid[0:a.shape[0],0:columnIdx.size-1][0].flatten()
b[rowIndexes,columnIndexes] = b[rowIndexes,columnIndexes]+1

B 现在应该是 array([[1,1,2],[1,1,2]])，因为它只针对 a 的列中的最大值的索引对 b 执行了操作。

有人知道更好的方法吗？最好只使用布尔掩码数组，以便我可以将此代码移植到 GPU 上运行，而无需太多麻烦。谢谢！

4

1 回答 1

1

我会建议一个答案，但数据略有不同。

c = np.array([[0,1,1],[2,1,0]])  # note that this data has dupes for max in row 1
d = np.array([[0,10,10],[20,10,0]]) # data to be chaged
c_argmax = np.argmax(c,axis=1)[:,np.newaxis]
b_map1 = c_argmax == np.arange(c.shape[1])
# now use the bool map as you described
d[b_map1] += 1
d
[out]
array([[ 0, 11, 10],
       [21, 10,  0]])

请注意，我创建了一个具有最大数字副本的原件。上面的内容可以按照您的要求与 argmax 一起使用，但您可能希望增加所有最大值。如：

c_max = np.max(c,axis=1)[:,np.newaxis]
b_map2 = c_max == c
d[b_map2] += 1
d
[out]
array([[ 0, 12, 11],
       [22, 10,  0]])
于 2013-06-02T03:25:54.803 回答