使用
np.take与
axis关键字参数:
>>> a = np.arange(2*3*4).reshape(2, 3, 4)>>> aarray([[[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]])>>> b = np.arange(3)>>> np.random.shuffle(b)>>> barray([1, 0, 2])>>> np.take(a, b, axis=1)array([[[ 4, 5, 6, 7], [ 0, 1, 2, 3], [ 8, 9, 10, 11]], [[16, 17, 18, 19], [12, 13, 14, 15], [20, 21, 22, 23]]])
如果要使用花式索引,则只需在索引元组中填充足够的空片即可:
>>> a[:, b]array([[[ 4, 5, 6, 7], [ 0, 1, 2, 3], [ 8, 9, 10, 11]], [[16, 17, 18, 19], [12, 13, 14, 15], [20, 21, 22, 23]]])
或更一般的情况下:
>>> axis = 1>>> idx = (slice(None),) * axis + (b,)>>> a[idx]array([[[ 4, 5, 6, 7], [ 0, 1, 2, 3], [ 8, 9, 10, 11]], [[16, 17, 18, 19], [12, 13, 14, 15], [20, 21, 22, 23]]])
但
np.take实际上应该是您的首选。



