如何在不改变形状的情况下对三维数组使用numpy lexsort？

问题背景 #

给定三维NumPy数组：

import numpy as np

arr = np.array([
[[1, 2, 0], [1, 0, 1], [1, 0, 0]],
[[1, 2, 1], [1, 0, 2], [1, 2, 0]],
[[1, 2, 2], [2, 0, 0], [1, 0, 3]]
])

需要对每个第一维度元素（即每个arr[i]）内部的长度为3的子数组，按以下优先级排序：

1. 第一排序键：子数组的最后一个元素（索引2，升序）
2. 第二排序键：子数组的第二个元素（索引1，升序）
3. 第三排序键：子数组的第一个元素（索引0，升序）

预期结果：

arr = np.array([
[[1, 0, 0], [1, 2, 0], [1, 0, 1]],
[[1, 2, 0], [1, 2, 1], [1, 0, 2]],
[[2, 0, 0], [1, 2, 2], [1, 0, 3]]
])

错误原因分析 #

你之前直接使用np.lexsort时未指定排序轴，导致函数对整个数组全局排序，返回的一维索引用于三维数组索引时触发广播，最终数组形状异常。np.lexsort默认会将数组展平后处理，而非针对每个子数组独立排序。

矢量化解决方案（无循环） #

利用np.lexsort的axis参数实现对每个子数组的独立排序，完全矢量化，适合处理大内存数组：

import numpy as np

arr = np.array([
[[1, 2, 0], [1, 0, 1], [1, 0, 0]],
[[1, 2, 1], [1, 0, 2], [1, 2, 0]],
[[1, 2, 2], [2, 0, 0], [1, 0, 3]]
])

# 按 lexsort 要求，排序键参数顺序为：最次键 → 主键
keys = (arr[..., 0], arr[..., 1], arr[..., 2])
# 在轴1（每个第一维度元素的内部数组）上执行 lexsort，得到排序索引
sorted_indices = np.lexsort(keys, axis=1)
# 利用索引对原数组排序，保持形状不变
result = arr[np.arange(arr.shape[0])[:, None], sorted_indices]

print(result)

代码解释 #

1. 排序键定义：np.lexsort的参数顺序是从最次键到主键，因此传入(arr[...,0], arr[...,1], arr[...,2])，对应第三、第二、第一排序键。
2. 指定排序轴：axis=1表示对每个第一维度元素（arr[i]）的内部数组（轴1）独立排序，返回的sorted_indices形状为(3,3)，每行对应原数组该行的排序索引。
3. 索引匹配：np.arange(arr.shape[0])[:, None]生成形状为(3,1)的索引数组，与sorted_indices广播后，为每个第一维度元素匹配对应的排序索引，最终输出形状保持(3,3,3)。

效率说明 #

该方法完全基于NumPy矢量化操作，避免了Python循环的开销，能够高效处理超大规模数组（如80GB级别的数据），充分利用NumPy的底层优化。

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