如何对三维NumPy数组按指定规则进行重新排序？

原始数组 #

现有三维NumPy数组arr：

import numpy as np
arr = np.array([
    [[ 0,  9,  0,  2], [ 6,  8,  0,  2]],
    [[ 0,  9,  1,  5], [ 4,  8,  1,  5]],
    [[ 1,  0,  1,  2], [ 6,  6,  1,  2]],
    [[ 0,  9,  1,  5], [ 5,  5,  1,  5]],
    [[ 0,  9,  1,  5], [ 6,  5,  1,  5]],
    [[ 1,  0,  1,  2], [ 4,  7,  1,  2]],  
    [[ 0,  9,  3,  4], [ 6,  4,  3,  4]]
])

排序需求 #

创建同形状的新数组new_arr，按以下优先级对arr中的二维子数组进行字典序排序（子数组本身保持不变，仅调整在三维数组中的顺序）：

1. 优先按arr[:, 0, 0]的值排序
2. 其次按arr[:, 1, 0]的值排序
3. 接着按arr[:, 0, 1]的值排序
4. 最后按arr[:, 1, 1]的值排序

解决方案 #

使用NumPy的np.lexsort()方法处理多键排序，该方法会从传入的最后一个参数开始，依次按字典序排序。

代码实现 #

# 提取所有排序键
key1 = arr[:, 0, 0]
key2 = arr[:, 1, 0]
key3 = arr[:, 0, 1]
key4 = arr[:, 1, 1]

# 生成排序索引：lexsort参数顺序为【最次要键→最主要键】
sorted_indices = np.lexsort([key4, key3, key2, key1])

# 根据索引重新排列数组
new_arr = arr[sorted_indices]

最终结果 #

运行上述代码后，new_arr的结果如下：

print(new_arr)
# 输出：
[[[ 0  9  1  5]
  [ 4  8  1  5]]

 [[ 0  9  1  5]
  [ 5  5  1  5]]

 [[ 0  9  3  4]
  [ 6  4  3  4]]

 [[ 0  9  1  5]
  [ 6  5  1  5]]

 [[ 0  9  0  2]
  [ 6  8  0  2]]

 [[ 1  0  1  2]
  [ 4  7  1  2]]

 [[ 1  0  1  2]
  [ 6  6  1  2]]]

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