Transform one-dimensional Numpy array into mask suitable for array index update of 2-dimensional array

Question

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.

How can I transform b into the "mask" version below?

score 1 · Accepted Answer · answered Aug 23 '23 at 16:40

1

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]])

answered Aug 23 '23 at 16:40

mozway

194,879
13
39
75

Ah I see, this is what I tried doing with `a[:,b]` but did not work as I expected – Alex Aug 23 '23 at 16:48
note that without clearing `a`, (4,1) is still set to 1, though it's not in the example result! – ti7 Aug 23 '23 at 16:58
1

@ti7 it was already set to 1, the first operation only changes the assigned indices, If you need to build an array with only the 1s and all 0s for the rest use the `mask` ;) – mozway Aug 23 '23 at 17:02
indeed / just so! – ti7 Aug 23 '23 at 17:05
2

For my actual use case I needed to update only a subset of "a" - so went with your mask solution above which I was able to modify - worked fantastic – Alex Aug 23 '23 at 17:07

score 0 · Answer 2 · answered Aug 23 '23 at 16:53

A (probably worse) alternate method would be to compute the offsets for the unraveled array and use unravel_index, though this might be practical if you had more than 2 dimensions or to make some sparse matrix

>>> idx = b + np.arange(len(b))*a.shape[-1]
>>> idx
array([ 1,  6, 10, 12, 19, 23])
>>> idx_multi = np.unravel_index(idx, a.shape)
>>> idx_multi
(array([0, 1, 2, 3, 4, 5]), array([1, 2, 2, 0, 3, 3]))
>>> a[:] = 0  # clear a
>>> a[idx_multi] = 1
>>> a
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]])

Transform one-dimensional Numpy array into mask suitable for array index update of 2-dimensional array

2 Answers2