Numpy中连续索引操作的更简洁等价实现方法求助

给定3维NumPy数组：

import numpy as np

arr = np.array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],

       [[ 6,  7],
        [ 8,  9],
        [10, 11]],

       [[12, 13],
        [14, 15],
        [16, 17]]])

原代码通过两次索引得到目标结果：

idx0, idx1 = [0,1], [0,2]
result = arr[idx0,...][:,idx1,:]

返回结果：

array([[[ 0,  1],
        [ 4,  5]],

       [[ 6,  7],
        [10, 11]]])

更优的单次索引实现 #

可以利用np.ix_生成多维索引网格，直接完成单次索引，避免创建中间数组，更简洁高效：

idx0, idx1 = [0,1], [0,2]
result = arr[np.ix_(idx0, idx1), :]

或者通过手动扩展索引维度实现等价效果：

result = arr[idx0[:, np.newaxis], idx1, :]

说明 #

原代码的两次索引会先创建一个中间数组（arr[idx0,...]的输出），再对这个中间数组做二次索引。而上述单次索引方式直接从原数组中定位目标元素，既节省内存开销，又提升了执行效率。

内容的提问来源于stack exchange，提问作者Dom Byrne