For each row of a numpy array, set specific elements to n-CodePudding

I have a 2d array arr. Here's an example:

>> arr
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])

I also have a 2d array of indices indices (with indices.shape[0] == arr.shape[0], and indices.shape[1] <= arr.shape[1]). Here is an example:

>>> indices
array([[0, 1],
       [1, 2],
       [0, 2],
       [0, 2]])

Now, I want to set some elements of arr to -1. Specifically, on the first row of arr, the first and the second element should be set to -1 (because indices[0] == [0, 1]). On the second row of arr, the second and the third element should be set to -1 (because indices[1] == [1, 2]). And so on.

This would be the expected result:

array([[ -1, -1,  2],
       [  3, -1, -1],
       [ -1,  7, -1],
       [ -1, 10, -1]])

I tried to look for existing solutions but I haven't found any. Any suggestion?

CodePudding user response：

Broadcast assignment:

>>> arr[np.arange(4)[:, None], indices] = -1
>>> arr
array([[-1, -1,  2],
       [ 3, -1, -1],
       [-1,  7, -1],
       [-1, 10, -1]])

CodePudding user response：

You could use a for loop:

for arow, irow in zip(arr, indices):
    for i in irow:
        arow[i] = -1
print(arr)

Or, in one line:

new_arr = np.array([[-1 if i in irow else arow[i] for i in range(len(arow))] for arow, irow in zip(arr, indices)])
print(new_arr)

Output:

[[-1 -1  2]
 [ 3 -1 -1]
 [-1  7 -1]
 [-1 10 -1]]