如何利用2D数组的排序索引对3D numpy数组按指定轴排序？

你的核心需求是将sort_idxs的排序规则应用到B的第三轴（axis=2），保持第一维度的对应关系，同时保留B的第二维度全部元素。可以通过扩展索引数组的维度实现广播匹配，以下是两种实用方法：

方法1：使用np.take_along_axis（语义清晰，推荐） #

take_along_axis可直接根据指定轴的索引数组提取元素，只需将sort_idxs扩展一个维度，使其与B的维度对齐：

# 将sort_idxs从(N,3)扩展为(N,1,3)，实现与B(N,3,3)的广播匹配
expanded_idxs = sort_idxs[:, np.newaxis, :]
# 对B的axis=2应用索引完成排序
sorted_B = np.take_along_axis(B, expanded_idxs, axis=2)

方法2：手动组合索引（直观理解维度对应） #

利用numpy索引的广播特性，分别生成各维度的索引后组合：

# 生成第一维度的索引，确保每个N对应的索引一致
n_idx = np.arange(B.shape[0])[:, np.newaxis, np.newaxis]
# 生成第二维度的索引，保留全部元素
axis1_idx = np.arange(B.shape[1])[np.newaxis, :, np.newaxis]
# 扩展sort_idxs的维度
expanded_idxs = sort_idxs[:, np.newaxis, :]
# 组合索引得到排序后的B
sorted_B = B[n_idx, axis1_idx, expanded_idxs]

验证示例 #

以N=2的场景测试：

import numpy as np

A = np.array([[3,1,2], [5,4,6]])
B = np.array([
    [[1,2,3], [4,5,6], [7,8,9]],
    [[10,11,12], [13,14,15], [16,17,18]]
])

sort_idxs = np.argsort(A, axis=1)
# sort_idxs输出：[[1,2,0], [1,0,2]]

expanded_idxs = sort_idxs[:, np.newaxis, :]
sorted_B = np.take_along_axis(B, expanded_idxs, axis=2)
# sorted_B输出：
# [[[2,3,1], [5,6,4], [8,9,7]],
#  [[11,10,12], [14,13,15], [17,16,18]]]

可见每个N对应的B矩阵，第三轴都按照A的排序规则完成了排序。

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