CODEWITHSUNDEEP
pythonnumpy
Lukas S
Lukas S

Reputation: 3603

apply 2d mask to 1d array in numpy

I have a 1d array and a 2d mask and would like to apply the 2d mask to the 1d array. By which I mean I want the indices where the array is true pluged into the 1d array giving me a n x 2 array where n is the number of trues in my array. E.g.

numbers = np.array([ 111, 2298, 5820, 6078, 7250])
mask = np.array([[False, False, False, False, False],
                 [ True, False, False, False, False],
                 [False, False, False, False, False],
                 [False, False, False, False, False],
                 [ True,  True, False, False, False]])
magic_function(arr=numbers, mask=mask)

Now since the first true has coordinates [1,0] I want the first result to be [2298, 111] and so on

array([[2298,111],[7250, 111],[2298,7250]])

Does this or something similar already exist in numpy?

Upvotes: 1

Views: 419

Answers (2)

hpaulj
hpaulj

Reputation: 231665

A variation on the other answer (argwhere is just np.transpose(np.nonzero(mask))):

In [26]: I,J=np.nonzero(mask)    
In [27]: I,J
Out[27]: (array([1, 4, 4]), array([0, 0, 1]))
In [29]: np.array([numbers[I],numbers[J]])
Out[29]: 
array([[2298, 7250, 7250],
       [ 111,  111, 2298]])
In [30]: np.stack([numbers[I],numbers[J]], axis=1)
Out[30]: 
array([[2298,  111],
       [7250,  111],
       [7250, 2298]])

Upvotes: 1

akuiper
akuiper

Reputation: 215117

You can use numpy.argwhere to create the array of indices needed for indexing:

numbers[np.argwhere(mask)]

#[[2298  111]
# [7250  111]
# [7250 2298]]

Upvotes: 4

Related Questions