通过Argmax沿另一个数组的轴遮罩一个2D numpy数组

Question

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

正确的方法是使用2d数组索引，但是它速度慢且令人费解，尤其是生成行索引的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列的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]])