Replace first occurence of 0 in a 2D numpy array

Question

I have a 2D array:

[[1,2,0,0],
[4,0,9,4],
[0,0,1,0],
[4,6,9,0]]

is there an efficient way (without using loops) to replace every first 0 in the array, with a 1:

[[1,2,1,0],
[4,1,9,4],
[1,0,1,0],
[4,6,9,1]]

?

Thanks a lot !

Are you working with a `numpy.ndarray`, or a *list* which is what you have posted above? — juanpa.arrivillaga, Jan 13 '22 at 15:47
Hi, what is it that you tried for this as code? e.g., can you detect 0s? — , Jan 13 '22 at 15:51

score 1 · Answer 1 · answered Jan 13 '22 at 15:54

So, you can use np.where to get the indices of the rows and columns where the array is 0:

In [45]: arr = np.array(
    ...:    [[1,2,0,0],
    ...:     [4,0,9,4],
    ...:     [0,0,1,0],
    ...:     [4,6,9,0]]
    ...: )

In [46]: r, c = np.where(arr == 0)

Then, use np.unique to get the unique x values, which will correspond to the first incidence of 0 in each row, and use return_index to get the indices to extract the corresponding column values:

In [47]: uniq_val, uniq_idx = np.unique(r, return_index=True)

In [48]: arr[uniq_val, c[uniq_idx]] = 1

In [49]: arr
Out[49]:
array([[1, 2, 1, 0],
       [4, 1, 9, 4],
       [1, 0, 1, 0],
       [4, 6, 9, 1]])

If performance is really an issue, you could just write a numba function, I suspect this would be very amenable to numba

rdesparbes · Accepted Answer · 2022-01-13T16:09:34.487

1

Here is a one-liner inspired by the accepted answer of this question:

a = np.array([
    [1, 2, 0, 0],
    [4, 0, 9, 4],
    [0, 0, 1, 0],
    [4, 6, 9, 0]
])
a[range(len(a)), np.argmax(a == 0, axis=1)] = 1

edited Jan 13 '22 at 16:09

answered Jan 13 '22 at 15:58

rdesparbes

75
5

Risks changing a nonzero entry in the array in case no 0 exists in a row. – Jan 13 '22 at 18:41

Replace first occurence of 0 in a 2D numpy array

2 Answers2