如何对3D numpy数组内各2D子数组按各自指定行单独排序列

问题原因 #

你原代码的错误在于a[:, 1, :].argsort()得到的是形状为(N, 3)的索引数组（N是3D数组包含的2D子数组数量），直接在最后一个维度传入该索引时，numpy会触发跨子数组的广播逻辑，无法实现「每个2D子数组单独使用自身的排序索引调整列顺序」的需求。

正确实现 #

你可以通过numpy高级索引匹配维度的方式实现需求，代码如下：

import numpy as np

# 示例3D数组
a = np.array([[[1,4,7],
               [3,1,5],
               [9,5,8]],
              [[2,8,7],
               [3,8,1],
               [9,2,8]]])

# 1. 按行方向，获取每个2D子数组第1行的列排序索引
sort_idx = a[:, 1, :].argsort(axis=1)
# 2. 构造匹配的维度索引，实现每个子数组单独排序
a = a[np.arange(a.shape[0])[:, np.newaxis], :, sort_idx]

运行后输出和你给出的期望结果完全一致。

逻辑说明 #

np.arange(a.shape[0])[:, np.newaxis]会生成形状为(N, 1)的第一维度索引，和形状为(N,3)的sort_idx广播后，会保证每个2D子数组只会调用自身对应的排序索引调整列顺序，不会出现跨子数组的索引混乱。

内容的提问来源于stack exchange，提问作者Christian S.