slicing a numpy array using indices from another array

Question

I am going crazy trying to slice a 2D array according to values from another array:

# array of integer values:
aa = np.random.randint(500, 600, size=(5,5)

array([[574, 550, 548, 545, 551],
       [547, 539, 539, 502, 528],
       [503, 530, 582, 567, 505],
       [590, 504, 510, 578, 525],
       [530, 548, 501, 580, 583]])

# array of indices:
ab = np.random.randint(4, size=(5,5))

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

What I want to return is a 2D subarray of aa, wherever ab is 1. But the closest I can get is:

aa[ab==1]

array([545, 551, 502, 582, 590, 578, 501, 580])

I'm always getting a 1D array as an output... How can I get my output array in the same dimensions as the original one?

EDIT: Sorry, I should have specified the expected output:

array([[545, 551],
       [502],
       [582],
       [590, 578],
       [501, 580]])

I also picked a poor example for the index array ab, it will always have the same number of 1s per row - so the output array would have the dimensions (5,2).

If it is the same dimensions as the original, what should be at the elements where ab!=1 do you just want to fill with 0? — Cory Kramer
– Cory Kramer, Commented Jun 9, 2020 at 14:33
If the suggestions from @CoryKramer is what you want, all you need is: numpy.where(ab == 1, aa, 0) — Paul H
– Paul H, Commented Jun 9, 2020 at 14:37
If you want a masked array, that's easy too, but you need to invert the logic: numpy.ma.masked_array(data=aa, mask=(ab!=1), fill_value=0) — Paul H
– Paul H, Commented Jun 9, 2020 at 14:38

Aleksander Lidtke · Accepted Answer · 2020-06-09 15:08:04Z

A 2D array of numbers of varying shape (one or two numbers per row) isn't, to the best of my knowledge and common sense, supported by numpy. If you try to instantiate the array that is your expected result, you'll hit issues:

In[23]: np.array([[545, 551],
   [502],
   [582],
   [590, 578],
   [501, 580]])
Out[23]: 
array([list([545, 551]), list([502]), list([582]), list([590, 578]),
   list([501, 580])], dtype=object) # Note the data type.

In[24]: np.array([[545, 551],
   [502],
   [582],
   [590, 578],
   [501, 580]],dtype=float) # Force float.
TypeError: float() argument must be a string or a number, not 'list'

If you're sure that you'll always have an equal and known number of 1s in ab, you could do the following:

In[2]: aa = np.array([[574, 550, 548, 545, 551],
   [547, 539, 539, 502, 528],
   [503, 530, 582, 567, 505],
   [590, 504, 510, 578, 525],
   [530, 548, 501, 580, 583]])

In[3]: ab = np.array([[3, 0, 2, 1, 1], # N.B. edited w.r.t. your example to have two 1s per row.
   [3, 1, 2, 1, 3],
   [0, 1, 1, 2, 0],
   [1, 2, 3, 1, 3],
   [3, 0, 1, 1, 0]])

In[4]: aa[ab==1].reshape((-1,2))
Out[4]: 
array([[545, 551],
   [539, 502],
   [530, 582],
   [590, 578],
   [501, 580]])

If you don't a priori know how many 1s there are in every row of ab, you could simply do the following:

In[5]: aa[ab==1].reshape((-1,sum(ab[0,:]==1)))
Out[5]: 
array([[545, 551],
   [539, 502],
   [530, 582],
   [590, 578],
   [501, 580]])

Does this answer your question? Or are you looking for something more like the MaskedArray suggested in comments?

This, especially the last part, is what I was looking for. Thanks @Aleksander for working out the reshape() trickery for me!

supreeth2812 · Accepted Answer · 2020-06-09 15:19:19Z

After a lot of trial and error, this is the best and shortest way I could come up with:

ac = np.array(list(zip(aa.ravel(),ab.ravel())), dtype=('i4,i4')).reshape(aa.shape)
ad= list(filter(None, [[i[0] for i in line if i[1] == 1] for line in ac]))
ad

Example:

aa = np.random.randint(500, 600, size=(5,5))

array([[560, 518, 569, 508, 582],
       [536, 556, 564, 523, 559],
       [501, 526, 596, 586, 567],
       [586, 536, 561, 503, 551],
       [500, 537, 508, 585, 556]])

ab = np.random.randint(4, size=(5,5))

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

#Zipping both arrays to create a tuple
ac = np.array(list(zip(aa.ravel(),ab.ravel())), dtype=('i4,i4')).reshape(aa.shape)
ac

array([[(560, 2), (518, 2), (569, 3), (508, 2), (582, 3)],
       [(536, 0), (556, 0), (564, 1), (523, 0), (559, 1)],
       [(501, 2), (526, 1), (596, 2), (586, 0), (567, 2)],
       [(586, 3), (536, 0), (561, 2), (503, 3), (551, 2)],
       [(500, 0), (537, 2), (508, 2), (585, 1), (556, 0)]])

#List comprehension to filter for elements where ab == 1 and remove empty lists
ad= list(filter(None, [[i[0] for i in line if i[1] == 1] for line in ac]))
ad

[[564, 559], [526], [585]]

Hope this was helpful!

hpaulj · Accepted Answer · 2020-06-09 15:55:41Z

The number of matches per row of ab vary:

In [246]: ab==1                                                                               
Out[246]: 
array([[False, False, False,  True,  True],
       [False, False, False,  True, False],
       [False, False,  True, False, False],
       [ True, False, False,  True, False],
       [False, False,  True,  True, False]])
In [247]: (ab==1).sum(axis=1)                                                                 
Out[247]: array([2, 1, 1, 2, 2])
In [248]: aa[ab==1]                                                                           
Out[248]: array([545, 551, 502, 582, 590, 578, 501, 580])

numpy indexing docs is quite clear that when the boolean indexing array matches the argument, the result is 1d. It can't in the general case maintain the original number of dimensions. In special cases you could reshape the result to 2d (but not this).

To get the ragged list of matches that you want, you have to iterate by rows:

In [249]: [a[b] for a,b in zip(aa, ab==1)]                                                    
Out[249]: 
[array([545, 551]),
 array([502]),
 array([582]),
 array([590, 578]),
 array([501, 580])]

https://numpy.org/doc/1.18/reference/arrays.indexing.html#boolean-array-indexing

If obj.ndim == x.ndim, x[obj] returns a 1-dimensional array filled with the elements of x corresponding to the True values of obj. The search order will be row-major, C-style.

Alternatives mentioned in the comments:

In [250]: np.where(ab==1, aa,0)                                                               
Out[250]: 
array([[  0,   0,   0, 545, 551],
       [  0,   0,   0, 502,   0],
       [  0,   0, 582,   0,   0],
       [590,   0,   0, 578,   0],
       [  0,   0, 501, 580,   0]])
In [251]: np.ma.masked_array(aa, ab!=1)                                                       
Out[251]: 
masked_array(
  data=[[--, --, --, 545, 551],
        [--, --, --, 502, --],
        [--, --, 582, --, --],
        [590, --, --, 578, --],
        [--, --, 501, 580, --]],
  mask=[[ True,  True,  True, False, False],
        [ True,  True,  True, False,  True],
        [ True,  True, False,  True,  True],
        [False,  True,  True, False,  True],
        [ True,  True, False, False,  True]],
  fill_value=999999)

The aa[ab==1] indexing is equivalent to indexing with the two arrays derived with nonzero (aka where):

In [252]: np.nonzero(ab==1)                                                                   
Out[252]: (array([0, 0, 1, 2, 3, 3, 4, 4]), array([3, 4, 3, 2, 0, 3, 2, 3]))
In [253]: aa[_]                                                                               
Out[253]: array([545, 551, 502, 582, 590, 578, 501, 580])

The shape of [253] matches the shape of the arrays in [252].

Collectives™ on Stack Overflow

slicing a numpy array using indices from another array

3 Answers 3

1 Comment

Comments

Comments

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

Comments

Comments

Related