Numpy update 2-dimensional array using array of indices

August 23, 2023

Given a 2-dimensional array a, I want to update select indices specified by b to a fixed value of 1.

test data:

import numpy as np

a = np.array(
    [[0, 1, 0, 0],
    [0, 0, 0, 0],
    [0, 0, 1, 0],
    [1, 0, 0, 0],
    [0, 1, 0, 1],
    [0, 0, 0, 0]]
)

b = np.array([1, 2, 2, 0, 3, 3])

One solution is to transform b into a masked array like this:

array([[0, 1, 0, 0],
       [0, 0, 1, 0],
       [0, 0, 1, 0],
       [1, 0, 0, 0],
       [0, 0, 0, 1],
       [0, 0, 0, 1]])

which would allow me to do a[b.astype(bool)] = 1 and solve the problem, but I’m not sure how to do that.

>Solution :

No need to build the mask, use indexing directly:

a[np.arange(len(b)), b] = 1

Output:

array([[0, 1, 0, 0],
       [0, 0, 1, 0],
       [0, 0, 1, 0],
       [1, 0, 0, 0],
       [0, 1, 0, 1],
       [0, 0, 0, 1]])

That said, the mask could be built using:

mask = b[:,None] == np.arange(a.shape[1])

Output:

array([[False,  True, False, False],
       [False, False,  True, False],
       [False, False,  True, False],
       [ True, False, False, False],
       [False, False, False,  True],
       [False, False, False,  True]])