替换 2D numpy 数组中第一次出现的 0

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 !

Answer 1

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

Answer 2

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

Answer 3

Im not sure if this is that efficent, but its better than loops at least. We search for the first 0 in each row, then extend this to coordinates and set those to 1.

filter = np.argmax(a == 0, axis=1)
f = np.column_stack((np.array(range(4)),filter))
a[[*f.T]] = 1