如何用Numpy高效实现2D数组按指定轴的另一个2D索引数组重排？

嘿，我来帮你搞定这个数组重排的问题！你想要用argsort得到的索引数组来重新排列原数组的每一行，其实Numpy有个专门的函数能完美解决这个需求——np.take_along_axis，它既高效又直观。

具体实现步骤 #

首先先重现你的初始数组和索引数组：

import numpy as np

in_arr = np.array([[1,2],[4,3]])
col_sort = np.argsort(in_arr, axis=1)

接下来用take_along_axis完成重排：

reordered_in_arr = np.take_along_axis(in_arr, col_sort, axis=1)

执行后得到的结果就是你想要的：

array([[1, 2],
       [3, 4]])

为什么这个方法高效？ #

take_along_axis是Numpy原生的向量化操作，完全利用了Numpy的底层优化，比手动循环或者自定义索引切片快得多，尤其是处理大规模数组的时候优势更明显。而且它的参数设计非常贴合“沿指定轴按索引选取元素”的场景，代码可读性拉满，后续做复杂按列操作也能无缝衔接。

补充：另一种高级索引方式 #

如果你好奇其他实现思路，也可以用Numpy的高级索引来实现：

# 生成每行对应的索引
row_indices = np.arange(in_arr.shape[0])[:, np.newaxis]
reordered_in_arr = in_arr[row_indices, col_sort]

不过这种方式需要手动构造行索引，相比之下take_along_axis更简洁省心，推荐优先使用。

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