python - numpy如何将数组变为除最大值之外的零一

Question

我想创建一个函数，当输入一个数组时，它返回一个形状相同但全为零的数组，期望 1 个值是最大的值。例如。像这样的数组:

my_array = np.arange(9).reshape((3,3))

[[ 0.  1.  2.]
 [ 3.  4.  5.]
 [ 6.  7.  8.]]

当传入函数时，我希望它像这样输出:

[[ 0.  0.  0.]
 [ 0.  0.  0.]
 [ 0.  0.  8.]]

异常(exception):

当有许多相等的最大值时，我只想要其中的一个，其余的被清零(顺序无关紧要)。

老实说，我不知道如何以一种优雅且高效的方式做到这一点，你会怎么做？

Answer 1

为了提高效率，请使用array-initialization和argmax来获取最大索引(如果多个索引，则第一个线性索引)-

def app_flat(my_array):
    out = np.zeros_like(my_array)
    idx  = my_array.argmax()
    out.flat[idx] = my_array.flat[idx]
    return out

我们还可以使用 ndarray.ravel() 代替 ndaarray.flat，我认为性能数字是相当的。

对于这种稀疏输出，为了获得内存效率和性能，您可能需要使用稀疏矩阵，尤其是对于大型数组。因此，对于稀疏矩阵输出，我们会有另一种选择，如下所示 -

from scipy.sparse import coo_matrix

def app_sparse(my_array):
    idx  = my_array.argmax()
    r,c = np.unravel_index(idx, my_array.shape)
    return coo_matrix(([my_array[r,c]],([r],[c])),shape=my_array.shape)

示例运行 -

In [336]: my_array
Out[336]: 
array([[0, 1, 2],
       [3, 4, 5],
       [8, 7, 8]])

In [337]: app_flat(my_array)
Out[337]: 
array([[0, 0, 0],
       [0, 0, 0],
       [8, 0, 0]])

In [338]: app_sparse(my_array)
Out[338]: 
<3x3 sparse matrix of type '<type 'numpy.int64'>'
    with 1 stored elements in COOrdinate format>

In [339]: app_sparse(my_array).toarray() # just to confirm values
Out[339]: 
array([[0, 0, 0],
       [0, 0, 0],
       [8, 0, 0]])

更大数组上的运行时测试 -

In [340]: my_array =  np.random.randint(0,1000,(5000,5000))

In [341]: %timeit app_flat(my_array)
10 loops, best of 3: 34.9 ms per loop

In [342]: %timeit app_sparse(my_array) # sparse matrix output
100 loops, best of 3: 17.2 ms per loop