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]

其中前2个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限制的地方。

Answer 2

它并没有完全解决您的问题，但choose()应该有所帮助：

>>> A = array(range(1, 28)).reshape(3, 3, 3)
>>> B = array([0, 0, 0, 1, 1, 1, 2, 2, 2]).reshape(3, 3)
>>> B.choose(A)
array([[ 1,  2,  3],
       [13, 14, 15],
       [25, 26, 27]])

它选择第一个维度而不是最后一个维度。

如何沿着nD数组的轴选择具有该轴的（n-1）D索引数组的值？

2 个答案: