如何按指定行对NumPy三维数组排序并同步关联行位置

问题描述 #

现有形状为2×3×4的NumPy数组，内容如下：

import numpy as np
arr = np.array([
    [[1,3,2,4],
     [0,1,0,1],
     [0,1,2,3]],
    [[1,4,2,3],
     [1,0,1,1],
     [0,1,2,3]]
])

需要对每个二维子矩阵单独排序：以子矩阵第一行的升序排列规则为基准，第二、第三行随第一行的索引变动同步移动对应位置元素，最终得到目标输出：

[[[1 2 3 4]
  [0 0 1 1]
  [0 2 1 3]]
 [[1 2 3 4]
  [1 1 1 0]
  [0 2 3 1]]]

实现思路 #

核心逻辑是先获取每个子矩阵第一行排序后的列位置索引，再通过NumPy高级索引用该索引重排所有子矩阵的列，保证同一列的三行元素同步位移：

• 第一步：切取所有子矩阵的第一行，调用argsort()得到每个子矩阵第一行升序对应的原列索引
• 第二步：构造匹配维度的批次索引，结合列索引对数组做高级索引，一次性完成所有子矩阵的列重排

可直接运行的代码 #

import numpy as np

# 原始输入数组
arr = np.array([
    [[1,3,2,4],
     [0,1,0,1],
     [0,1,2,3]],
    [[1,4,2,3],
     [1,0,1,1],
     [0,1,2,3]]
])

# 计算每个子矩阵第一行排序后的列索引，形状为(2,4)
sort_col_idx = arr[:, 0, :].argsort(axis=1)

# 构造批次索引配合高级索引
batch_idx = np.arange(arr.shape[0])[:, None]
# 按列索引重排，所有行同步跟随移动
sorted_arr = arr[batch_idx, :, sort_col_idx]

print(sorted_arr)

运行输出和目标结果完全一致：

[[[1 2 3 4]
  [0 0 1 1]
  [0 2 1 3]]

 [[1 2 3 4]
  [1 1 1 0]
  [0 2 3 1]]]

补充说明 #

• arr[:,0,:]切取的是所有二维子矩阵的第一行，argsort(axis=1)逐行计算升序排列对应的原位置索引，不会修改原数组
• 高级索引时batch_idx形状为(2,1)，sort_col_idx形状为(2,4)，广播后会对每个子矩阵的所有行，按对应列顺序取数，天然实现第二、第三行跟随第一行排序同步移动
• 该方法可直接套用到结构为(n, m, k)的同类型数组（n个m行k列的子矩阵），无需修改逻辑

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