CODEWITHSUNDEEP
pythonarraysnumpy
user2503169
user2503169

Reputation: 147

Return rows and columns from a 2d array using values from a 1d array in Numpy

I am trying to use the values of a 1-dim array to slice/return the rows and columns from a 2-dim array in Numpy. For example, say I have the following one dim array: [1,3,5)] and the following 2 dim array:

array([[1, 0, 0, 0, 0, 0],
       [0, 4, 0, 0, 0, 1],
       [0, 0, 3, 0, 0, 0],
       [0, 1, 0, 7, 0, 10],
       [0, 0, 0, 0, 8, 0],
       [0, 2, 0, 0, 0, 9]])

How do I return the following:

array([[4, 0, 1],
       [1, 7, 10],
       [2, 0, 9]])

I would also like to be able to produce a 6x6 mask using the same example. So that I would get this:

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

I have tried many different things and nothing seems to get exactly what I need. I know I could do it by writing a couple of loops, but I figured there must be an easier way. I have also done a number of searches and still no luck. Thanks in advance!

Upvotes: 1

Views: 1967

Answers (4)

Pablo
Pablo

Reputation: 2481

Maybe could be useful the routine np.meshgrid():

a = array([[1, 0, 0, 0, 0, 0],
   [0, 4, 0, 0, 0, 1],
   [0, 0, 3, 0, 0, 0],
   [0, 1, 0, 7, 0, 10],
   [0, 0, 0, 0, 8, 0],
   [0, 2, 0, 0, 0, 9]])

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

B = np.meshgrid(b,b)

print a[B].T

Out: [[ 4  0  1]
      [ 1  7 10]
      [ 2  0  9]]

I think that is the desired result.

Upvotes: 1

Alex Szatmary
Alex Szatmary

Reputation: 3571

mgilson's answer is what you want if you know that the strides are regular. However, if a is your array, and you want to get a[i, j] for all i and j in some 1-D array, b, there is a gotcha.

In : a = array([[1, 0, 0, 0, 0, 0],
       [0, 4, 0, 0, 0, 1],
       [0, 0, 3, 0, 0, 0],
       [0, 1, 0, 7, 0, 10],
       [0, 0, 0, 0, 8, 0],
       [0, 2, 0, 0, 0, 9]])

In : b = np.array([1, 3, 5])

In : a[b, b]
Out: array([4, 7, 9])

The problem is that this just gets you elements (1, 1), (3, 3) and (5, 5). If you try doing the indexing in separate steps,

In : a[b][b]
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-87-c3c286c42537> in <module>()
----> 1 a[b][b]

IndexError: index 3 is out of bounds for axis 0 with size 3

This error arises because I'm indexing the same axis twice. Either of these work:

In : a[b][:, b]
Out:
array([[ 4,  0,  1],
       [ 1,  7, 10],
       [ 2,  0,  9]])

In : a[:, b][b]
Out:
array([[ 4,  0,  1],
       [ 1,  7, 10],
       [ 2,  0,  9]])

The [b] picks out the rows, and [:, b] picks out the columns; these indexing operations can be swapped.

Upvotes: 0

mgilson
mgilson

Reputation: 309929

Is this what you want?

>>> a = array([[1, 0, 0, 0, 0, 0],
...        [0, 4, 0, 0, 0, 1],
...        [0, 0, 3, 0, 0, 0],
...        [0, 1, 0, 7, 0, 10],
...        [0, 0, 0, 0, 8, 0],
...        [0, 2, 0, 0, 0, 9]])
>>> 
>>> a[1::2,1::2]
array([[ 4,  0,  1],
       [ 1,  7, 10],
       [ 2,  0,  9]])

Since your stride access is so regular, you can accomplish this with basic slicing. As for the mask:

>>> a = np.ones(a.shape,dtype=bool)
>>> a[1::2,1::2] = False
>>> a
array([[ True,  True,  True,  True,  True,  True],
       [ True, False,  True, False,  True, False],
       [ True,  True,  True,  True,  True,  True],
       [ True, False,  True, False,  True, False],
       [ True,  True,  True,  True,  True,  True],
       [ True, False,  True, False,  True, False]], dtype=bool)

Of course, this answer is assuming you want every other element along the axis (starting with index 1). You could modify the slice to stop when the index is 6: a[1:6:2,1:6:2] or to take every 3rd element, a[1::3,1::3], but if you need random access into the array, that becomes a bit harder...

You can do something like this:

>>> b = [1,3,5]
>>> a[:,b][b]
array([[ 4,  0,  1],
       [ 1,  7, 10],
       [ 2,  0,  9]])
>>> a[b][:,b]  #I think the same thing, but depending on data layout, one may be faster than the other
array([[ 4,  0,  1],
       [ 1,  7, 10],
       [ 2,  0,  9]])

At this point though, you're probably making a copy of the array rather than just getting a view. This is less efficient and you won't be able to use it to construct the boolean mask as we did previously I don't think.

Upvotes: 5

Blckknght
Blckknght

Reputation: 104712

If you put your big array in a and your small one in i, you can do the slicing with a[i][...,i]. There's probably a way to do it in one step too, but I'm not enough of a numpy wizard to know it yet.

Upvotes: 0

Related Questions