Numpy：用零填充 ndarray 的非最大元素

Question

我有一个ndarray，我想把最后一个维度的所有非最大元素都设置为零。

a = np.array([[[1,8,3,4],[6,7,10,6],[11,12,15,4]],
              [[4,2,3,4],[4,7,9,8],[41,14,15,3]],
              [[4,22,3,4],[16,7,9,8],[41,12,15,43]]
             ])
print(a.shape)
(3,3,4)

我可以通过np.argmax()得到最大元素的索引:

b = np.argmax(a, axis=2)
b
array([[1, 2, 2],
       [0, 2, 0],
       [1, 0, 3]])

显然，b比a小1维。现在，我想要一个新的 3-d 数组，除了最大值所在的位置外，它全为零。

我想得到这个数组：

np.array([[[0,1,0,0],[0,0,1,0],[0,0,1,0]],
          [[1,0,0,1],[0,0,1,0],[1,0,0,0]],
          [[0,1,0,0],[1,0,0,0],[0,0,0,1]]
         ])

一种实现方式，我尝试创建这些临时数组

b = np.repeat(b[:,:,np.newaxis], 4, axis=2)
t = np.repeat(np.arange(4).reshape(4,1), 9, axis=1).T.reshape(b.shape)

z = np.zeros(shape=a.shape, dtype=int)
z[t == b] = 1
z
array([[[0, 1, 0, 0],
    [0, 0, 1, 0],
    [0, 0, 1, 0]],

   [[1, 0, 0, 0],
    [0, 0, 1, 0],
    [1, 0, 0, 0]],

   [[0, 1, 0, 0],
    [1, 0, 0, 0],
    [0, 0, 0, 1]]])

知道如何以更有效的方式获得它吗？

Answer 1

这是一种使用广播的方法：

In [108]: (a == a.max(axis=2, keepdims=True)).astype(int)
Out[108]: 
array([[[0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 1, 0]],

       [[1, 0, 0, 1],
        [0, 0, 1, 0],
        [1, 0, 0, 0]],

       [[0, 1, 0, 0],
        [1, 0, 0, 0],
        [0, 0, 0, 1]]])

Numpy：用零填充 ndarray 的非最大元素

Numpy: filling the non-maximum elements of ndarray with zeros

python

arrays

numpy

argmax