CODEWITHSUNDEEP
pythonarraysperformancenumpy
anon_swe
anon_swe

Reputation: 9335

Keep First Non-Zero Element, Set All Others to 0

I have a 2-d NumPy array that looks like this:

array([[0. , 0. , 0.2, 0.2],
       [0.3, 0. , 0.3, 0. ]])

I'd like to modify it so that each row consists of all 0's, except for the first non-zero entry. If it's all 0s to start with, we don't change anything.

I could do this:

example = np.array([[0,0, 0.2, 0.2], [0.3, 0, 0.3, 0]])
my_copy = np.zeros_like(example)
for i, row in enumerate(example):
    for j, elem in enumerate(row):
        if elem > 0:
            my_copy[i, j] = elem
            break

But that's ugly and not vectorized. Any suggestions for how to vectorize this?

Thanks!

Upvotes: 1

Views: 471

Answers (3)

user3483203
user3483203

Reputation: 51165

Setup

x = np.array([[0. , 0. , 0.2, 0.2],
          [0.3, 0. , 0.3, 0. ],
          [0. , 0. , 0. , 0. ]])

Using logical_and with np.eye:

m = (x!=0).argmax(1)
x[~np.logical_and(x, np.eye(x.shape[1])[m])] = 0

Output:

array([[0. , 0. , 0.2, 0. ],
       [0.3, 0. , 0. , 0. ],
       [0. , 0. , 0. , 0. ]])

Using this method will be slightly slower than the other two suggested.

Upvotes: 2

jpp
jpp

Reputation: 164623

Here's a vectorised solution. The trick is to calculate your first non-zero entries via bool conversion and argmax.

import numpy as np

A = np.array([[0. , 0. , 0.2, 0.2],
              [0.3, 0. , 0.3, 0. ],
              [0. , 0. , 0. , 0. ]])

res = np.zeros(A.shape)
idx =  np.arange(res.shape[0])
args = A.astype(bool).argmax(1)
res[idx, args] = A[idx, args]

print(res)

array([[ 0. ,  0. ,  0.2,  0. ],
       [ 0.3,  0. ,  0. ,  0. ],
       [ 0. ,  0. ,  0. ,  0. ]])

Upvotes: 5

rafaelc
rafaelc

Reputation: 59274

Simply

e =np.zeros(example.shape)
rows = np.arange(example.shape[0])
cols =  np.argmax(example != 0, 1)
e[rows, cols] = example[rows, cols]

Upvotes: 3

Related Questions