Python 通过Argmax沿另一个阵列的轴遮罩一个二维Numpy阵列

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

我使用2d数组索引的方法是正确的，但是它很慢而且有点复杂，特别是mgrid黑客生成行索引。对于这个例子，它只是[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]]），因为它只对B执行沿a列的max索引的操作

---

有人知道更好的方法吗？最好只使用布尔掩蔽数组，这样我就可以将代码移植到GPU上运行，而不会有太多麻烦。谢谢

我会给出一个答案，但数据略有不同

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，但您可能希望增加所有max值。例如：

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

---

我将提出一个答案，但数据略有不同

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，但您可能希望增加所有max值。例如：

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