问题描述
这是由激发的.
给出形状为(n0，n1)的数组A和形状为(n0)的数组J，我想创建形状为(n0)的数组B，以使
Given array A with shape (n0,n1), and array J with shape (n0), I'd like to create an array B with shape (n0) such that
B[i] = A[i,J[i]]
我也想将其推广到k维数组，其中A的形状为(n0,n1,...,nk)，J的形状为(n0,n1,...,n(k-1))
I'd also like to be able to generalize this to k-dimensional arrays, where A has shape (n0,n1,...,nk) and J has shape (n0,n1,...,n(k-1))
有一些混乱的，扁平化的方法可以对索引顺序进行假设:
There are messy, flattening ways of doing this that make assumptions about index order:
import numpy as np
B = A.ravel()[ J+A.shape[-1]*np.arange(0,np.prod(J.shape)).reshape(J.shape) ]
问题是，有没有一种方法不依赖于展平数组和手动处理索引?
The question is, is there a way to do this that doesn't rely on flattening arrays and dealing with indexes manually?
推荐答案
对于2d和1d情况，此索引有效:
For the 2 and 1d case, this indexing works:
A[np.arange(J.shape[0]), J]
通过重塑为2d(并返回)可以将其应用于更多尺寸:
Which can be applied to more dimensions by reshaping to 2d (and back):
A.reshape(-1, A.shape[-1])[np.arange(np.prod(A.shape[:-1])).reshape(J.shape), J]
对于3d A，此方法有效:
For 3d A this works:
A[np.arange(J.shape[0])[:,None], np.arange(J.shape[1])[None,:], J]
其中第1 2个arange索引的广播尺寸与J相同.
where the 1st 2 arange indices broadcast to the same dimension as J.
对于lib.index_tricks中的函数，可以表示为:
With functions in lib.index_tricks, this can be expressed as:
A[np.ogrid[0:J.shape[0],0:J.shape[1]]+[J]]
A[np.ogrid[slice(J.shape[0]),slice(J.shape[1])]+[J]]
或用于多个维度:
A[np.ix_(*[np.arange(x) for x in J.shape])+(J,)]
A[np.ogrid[[slice(k) for k in J.shape]]+[J]]
对于较小的A和J(例如2 * 3 * 4)，J.choose(np.rollaxis(A,-1))更快.所有额外的时间都在准备索引元组中. np.ix_比np.ogrid快.
For small A and J (eg 2*3*4), J.choose(np.rollaxis(A,-1)) is faster. All of the extra time is in preparing the index tuple. np.ix_ is faster than np.ogrid.
np.choose具有大小限制.在上端，它比ix_慢:
np.choose has a size limit. At its upper end it is slower than ix_:
In [610]: Abig=np.arange(31*31).reshape(31,31)
In [611]: Jbig=np.arange(31)
In [612]: Jbig.choose(np.rollaxis(Abig,-1))
Out[612]:
array([ 0, 32, 64, 96, 128, 160, ... 960])
In [613]: timeit Jbig.choose(np.rollaxis(Abig,-1))
10000 loops, best of 3: 73.1 µs per loop
In [614]: timeit Abig[np.ix_(*[np.arange(x) for x in Jbig.shape])+(Jbig,)]
10000 loops, best of 3: 22.7 µs per loop
In [635]: timeit Abig.ravel()[Jbig+Abig.shape[-1]*np.arange(0,np.prod(Jbig.shape)).reshape(Jbig.shape) ]
10000 loops, best of 3: 44.8 µs per loop
我在 https://stackoverflow.com/a/28007256/901925 上做了类似的索引测试，发现对于更大的数组(例如n0=1000)，flat索引编制更快.那是我了解choice的32个限制的地方.
I did similar indexing tests at https://stackoverflow.com/a/28007256/901925, and found that flat indexing was faster for much larger arrays (e.g. n0=1000). That's where I learned about the 32 limit for choice.
这篇关于如何选择具有(n-1)D个轴索引的nD阵列轴上的值?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！