做这件事有很多种方法。我们首先生成一些测试数据:
In [1]: a = np.random.rand(10, 10, 20)
In [2]: b = np.random.randint(20, size=(10,10)) # random integers in range 0..19
解决这个问题的一种方法是创建两个索引向量,其中一个是行向量,另一个是 0..9 的列向量,使用meshgrid http://docs.scipy.org/doc/numpy/reference/generated/numpy.meshgrid.html:
In [3]: i1, i0 = np.meshgrid(range(10), range(10), sparse=True)
In [4]: c = a[i0, i1, b]
这有效是因为i0
, i1
and b
将全部广播到 10x10 矩阵。快速测试正确性:
In [5]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10))
Out[5]: True
另一种方法是使用choose http://docs.scipy.org/doc/numpy/reference/generated/numpy.choose.html and rollaxis http://docs.scipy.org/doc/numpy/reference/generated/numpy.rollaxis.html:
# choose needs a sequence of length 20, so move last axis to front
In [22]: aa = np.rollaxis(a, -1)
In [23]: c = np.choose(b, aa)
In [24]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10))
Out[24]: True