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

时间:2013-06-01 20:14:48

标签: python numpy masking

我有一个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上而不会有太多麻烦。谢谢!

1 个答案:

答案 0 :(得分: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]])