CODEWITHSUNDEEP
pythonarraysnumpy
Tissuebox
Tissuebox

Reputation: 1056

numpy how to change an array into a zero one except the max value

I want to make a function that when fed an array, it returns an array of the same shape but with all zeros expect for 1 value that is the max one. eg. with an array like this:

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

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

when passed in the function I want it out like this:

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

exeption:

When there is many max value that are equal, I only want one out of them and the rest gets zero'ed out (the order doesn't matter).

I am honeslty clueless as to how to make this in an elegant way that is furthermore efficient, how would you do it?

Upvotes: 2

Views: 1485

Answers (2)

Marco Cerliani
Marco Cerliani

Reputation: 22031

with few lines:

my_array = np.arange(9).reshape((3,3))
my_array2 = np.zeros(len(my_array.ravel()))
my_array2[np.argmax(my_array)] = np.max(my_array)
my_array2 = my_array2.reshape(my_array.shape)

Upvotes: 0

Divakar
Divakar

Reputation: 221594

For efficiency, use array-initialization and argmax to get the max index (first one linearly indexed if more than one) -

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

We can also use ndarray.ravel() in place of ndaarray.flat and I would think that the performance numbers would be comparable.

For this sparsey output, to gain memory efficiency and hence performance, you might want to use sparse matrices, especially for large arrays. Thus, for sparse matrix output, we would have an alternative one, like so -

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)

Sample run -

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

Runtime test on bigger array -

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

Upvotes: 2

Related Questions