如何用二维整数数组索引二维Numpy数组？

你遇到的问题是：用argsort得到每列的排序索引数组s后，直接用x[s]会得到(10,5,5)的三维数组，但你想要的是(10,5)的结果——也就是对x的每一列，用s对应列的索引来提取元素。

下面是几种高效的向量化实现方法，比列表推导式快得多：

方法1：配对行/列索引（最直观） #

生成和s匹配的列索引数组，和s一起作为索引：

import numpy as np
x = np.random.normal(size=(10, 5))
s = np.argsort(x, 0)

# 生成列索引，广播后和s的形状对齐
cols = np.arange(x.shape[1])
result = x[s, cols]

这里cols是[0,1,2,3,4]，NumPy会自动把它广播成(10,5)的形状，和s一一对应，最终取出每个位置(s[i,j], j)的元素，得到(10,5)的结果。

方法2：用take+对角线提取 #

如果习惯用take函数，也可以这么写：

result = np.take(x, s, axis=0).diagonal(axis1=1, axis2=2).T

np.take(x, s, axis=0)会得到和x[s]一样的三维数组，然后通过diagonal提取每一列对应的对角线元素，最后转置得到目标形状。

为什么直接x[s]不行？ #

当你用二维数组s索引x时，NumPy的规则是：s里的每个元素都是行索引，对于s[i,j]，它会取出x[s[i,j], :]这一整行，所以最终会堆叠成(10,5,5)的三维数组——相当于对s的每个位置都取了一整行，而不是你需要的单个元素。

而我们需要的是对每个(i,j)，取x[s[i,j], j]，所以必须明确指定列索引，让行和列一一对应。

效率对比 #

你用的列表推导式需要循环每一列，在数组规模大的时候，向量化方法的速度会是列表推导式的几倍甚至几十倍，完全利用了NumPy的底层优化。

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