Replace elements of numpy array based on first occurrence of a particular value

Question

Suppose there's a numpy 2D array as follows:

>>> x = np.array([[4,2,3,1,1], [1,0,3,2,1], [1,4,4,3,4]])
>>> x
array([[4, 2, 3, 1, 1],
       [1, 0, 3, 2, 1],
       [1, 4, 4, 3, 4]])

My objective is to - find the first occurrence of value 4 in each row, and set the rest of the elements (including that element) in that row to 0. Hence, after this operation, the transformed array should look like:

>>> x_new
array([[0, 0, 0, 0, 0],
       [1, 0, 3, 2, 1],
       [1, 0, 0, 0, 0]])

What is the pythonic and optimized way of doing this? I tried with a combination of np.argmax() and np.take() but was unable to achieve the end objective.

Dan · Accepted Answer · 2020-08-06 12:32:57Z

3

You can do it using a cumulative sum across the columns (i.e. axis=1) and boolean indexing:

n = 4
idx = np.cumsum(x == n, axis=1) > 0
x[idx] = 0

or maybe a better way is to do a cumulative (logical) or:

idx = np.logical_or.accumulate(x == n, axis=1)

edited Aug 6, 2020 at 12:32

answered Aug 6, 2020 at 12:27

Dan

45.8k20 gold badges98 silver badges169 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Replace elements of numpy array based on first occurrence of a particular value

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related