python - 如何选择沿 nD 数组轴的值以及该轴的 (n-1)D 索引数组？

Question

这是由 my answer here 激发的.

给定形状为 (n0,n1) 的数组 A 和形状为 (n0) 的数组 J，我想创建一个形状为 (n0) 的数组 B，这样

B[i] = A[i,J[i]]

我还希望能够将其推广到 k 维数组，其中 A 的形状为 (n0,n1,...,nk)，J 的形状为 ( n0,n1,...,n(k-1))

有一些凌乱的、扁平化的方法可以对索引顺序做出假设:

import numpy as np
B = A.ravel()[   J+A.shape[-1]*np.arange(0,np.prod(J.shape)).reshape(J.shape) ]

问题是，有没有一种方法可以做到这一点而不依赖于展平数组和手动处理索引？

Answer 1

对于 2 和 1d 情况，此索引有效:

A[np.arange(J.shape[0]), J]

可以通过整形为 2d(和返回)应用于更多维度:

A.reshape(-1, A.shape[-1])[np.arange(np.prod(A.shape[:-1])).reshape(J.shape), J]

对于 3d A 这有效:

A[np.arange(J.shape[0])[:,None], np.arange(J.shape[1])[None,:], J]

第一个和第二个 arange 索引广播到与 J 相同的维度。

使用 lib.index_tricks 中的函数，可以表示为:

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 更快。​​

np.choose 有大小限制。在其上端，它比 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 做了类似的索引测试，并发现 flat 索引对于更大的数组(例如 n0=1000)更快。这就是我了解 choice 的 32 限制的地方。