简体   繁体   中英

Indexing NumPy 2D array with another 2D array

I have something like

m = array([[1, 2],
            [4, 5],
            [7, 8],
            [6, 2]])

and

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

My target is

result = array([1, 5, 7, 6])

I tried _ix as I read at Simplfy row AND column extraction, numpy , but this did not result in what I wanted.

ps Please change the title of this question if you can think of a more precise one.

The numpy way to do this is by using np.choose or fancy indexing/take (see below):

m = array([[1, 2],
           [4, 5],
           [7, 8],
           [6, 2]])
select = array([0,1,0,0])

result = np.choose(select, m.T)

So there is no need for python loops, or anything, with all the speed advantages numpy gives you. mT is just needed because choose is really more a choise between the two arrays np.choose(select, (m[:,0], m[:1])) , but its straight forward to use it like this.


Using fancy indexing :

result = m[np.arange(len(select)), select]

And if speed is very important np.take , which works on a 1D view (its quite a bit faster for some reason, but maybe not for these tiny arrays):

result = m.take(select+np.arange(0, len(select) * m.shape[1], m.shape[1]))

I prefer to use NP.where for indexing tasks of this sort (rather than NP.ix_ )

What is not mentioned in the OP is whether the result is selected by location (row/col in the source array) or by some condition (eg, m >= 5). In any event, the code snippet below covers both scenarios.

Three steps:

  1. create the condition array ;

  2. generate an index array by calling NP.where , passing in this condition array; and

  3. apply this index array against the source array


>>> import numpy as NP

>>> cnd = (m==1) | (m==5) | (m==7) | (m==6)
>>> cnd
  matrix([[ True, False],
          [False,  True],
          [ True, False],
          [ True, False]], dtype=bool)

>>> # generate the index array/matrix 
>>> # by calling NP.where, passing in the condition (cnd)
>>> ndx = NP.where(cnd)
>>> ndx
  (matrix([[0, 1, 2, 3]]), matrix([[0, 1, 0, 0]]))

>>> # now apply it against the source array   
>>> m[ndx]
  matrix([[1, 5, 7, 6]])


The argument passed to NP.where, cnd , is a boolean array, which in this case, is the result from a single expression comprised of compound conditional expressions (first line above)

If constructing such a value filter doesn't apply to your particular use case, that's fine, you just need to generate the actual boolean matrix (the value of cnd ) some other way (or create it directly).

恕我直言,这是最简单的变体:

m[np.arange(4), select]

Since the title is referring to indexing a 2D array with another 2D array, the actual general numpy solution can be found here .

In short: A 2D array of indices of shape (n,m) with arbitrary large dimension m , named inds , is used to access elements of another 2D array of shape (n,k), named B :

# array of index offsets to be added to each row of inds
offset = np.arange(0, inds.size, inds.shape[1])

# numpy.take(B, C) "flattens" arrays B and C and selects elements from B based on indices in C
Result = np.take(B, offset[:,np.newaxis]+inds)

Another solution, which doesn't use np.take and I find more intuitive, is the following:

B[np.expand_dims(np.arange(B.shape[0]), -1), inds]

The advantage of this syntax is that it can be used both for reading elements from B based on inds (like np.take ), as well as for assignment.

使用python怎么样?

result = array([subarray[index] for subarray, index in zip(m, select)])
result = array([m[j][0] if i==0 else m[j][1] for i,j in zip(select, range(0, len(m)))])

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM