如何正确使用元组列表索引NumPy多维数组？

嘿，这个问题我太熟了——刚开始用NumPy处理多维索引的时候，几乎每个人都会踩这个坑！直接用a[t]得到的结果根本不是你想要的，而循环迭代又慢得让人抓狂，尤其是维度多、数据量大的时候。

先给你说清楚为什么a[t]不对：NumPy的索引规则是按列匹配维度，也就是说，t的第一列对应数组a的第一个维度，第二列对应第二个维度，以此类推。所以当你传入一个k行n列的t时，它会返回一个k维的数组（或对应形状的数组），而不是k个标量。

那正确的高效做法是什么呢？其实只需要把t转置一下，再转换成元组，直接传给a做索引就行：

def get(a, t):
    return a[tuple(t.T)]

举个实际的例子验证一下：

import numpy as np

# 构造一个3x3x3的测试数组
a = np.arange(27).reshape(3,3,3)
# 要取的3组多维索引
t = np.array([[0,1,2], [2,1,0], [1,0,1]])

# 你的循环方法结果
def get_loop(a, t):
    res = []
    for ind in t:
        s = a
        for i in ind:
            s = s[i]
        res.append(s)
    return np.array(res)

# 向量化方法结果
print(get_loop(a, t))  # 输出 [11 23  7]
print(a[tuple(t.T)])   # 输出 [11 23  7]

完全一致！而且这个方法是完全向量化的，底层用C实现，速度比循环快几个数量级——尤其是当你的t有几万甚至几十万行的时候，差距会特别明显。

另外还有一种等价的方法，用np.ravel_multi_index把多维索引转换成扁平的一维索引，再用np.take取值：

def get_ravel(a, t):
    flat_indices = np.ravel_multi_index(t.T, a.shape)
    return np.take(a, flat_indices)

这个方法的效率和前一种差不多，看你个人习惯选择。

为什么这两种方法快？因为NumPy的向量化操作避免了Python层面的循环开销，把所有计算都放到底层优化过的C代码里执行，充分利用了CPU的缓存和并行计算能力，而你的循环方法每次迭代都要在Python和NumPy之间切换，效率极低。

下次遇到类似的多维批量索引问题，记住这个技巧就行啦！

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