Question

我有一个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现在应该是数组（[[1,1,2]，[1,1,2]]），因为它只对a沿着a列的max的索引执行b上的操作。

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

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

通过沿另一个阵列的轴的Argmax掩盖一个2D Numpy阵列

1 个答案: