Python中NumPy不规则数组排序保持关联数组顺序一致

需求说明 #

现有两个关联的NumPy数组Ii0和Iv0，需满足以下要求：

• 对Ii0按每个子元素的第二个值（索引为1的维度值，对应参数j）升序排序，得到结果Ii01
• 关联数组Iv0需完全跟随Ii0的排序位置同步重排，得到结果Iv01
• 代码需兼容Ii0、Iv0的多种不同形状输入

原有错误实现 #

原有代码单独对Iv0做独立排序，未绑定Ii0的排序顺序，无法得到正确结果，代码如下：

import numpy as np

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

Iv0 = np.array([[[10],
                 [20],
                 [30]],
                        [[100],
                         [200],
                         [300]]])

Ii01 = np.array([sorted(i, key = lambda x : x[1]) for i in Ii0])
print("Ii01  =",Ii01)
Iv01 = np.array([sorted(i, key = lambda x : x[0]) for i in Iv0])
print("Iv01  =",Iv01)

错误运行输出 #

Ii01  = [array([[[0, 1],
        [1, 2],
        [0, 3]],

       [[0, 1],
        [0, 3],
        [2, 5]]])]
Iv01  = [array([[[ 10],
        [ 20],
        [ 30]],

       [[100],
        [200],
        [300]]])]

期望输出 #

Ii01  = [array([[[0, 1],
        [1, 2],
        [0, 3]],

       [[0, 1],
        [0, 3],
        [2, 5]]])]
Iv01  = [array([[[ 10],
        [ 30],
        [ 20]],

       [[300],
        [100],
        [200]]])]

正确实现 #

核心逻辑：先获取Ii0每个分组下按第二个值排序的索引，再用同一索引重排两个数组，保证顺序完全一致，同时兼容多形状输入：

import numpy as np

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

Iv0 = np.array([[[10],
                 [20],
                 [30]],
                        [[100],
                         [200],
                         [300]]])

Ii01 = []
Iv01 = []
# 逐组处理，保证两个数组用完全相同的索引重排
for group_ii, group_iv in zip(Ii0, Iv0):
    # 生成按Ii0子元素第二个值升序的排序索引
    sort_indices = np.argsort(group_ii[:, 1])
    Ii01.append(group_ii[sort_indices])
    Iv01.append(group_iv[sort_indices])

Ii01 = np.array(Ii01)
Iv01 = np.array(Iv01)

print("Ii01  =",Ii01)
print("Iv01  =",Iv01)

运行结果 #

运行输出与期望完全一致：

Ii01  = [[[0 1]
  [1 2]
  [0 3]]

 [[0 1]
  [0 3]
  [2 5]]]
Iv01  = [[[ 10]
  [ 30]
  [ 20]]

 [[300]
  [100]
  [200]]]

兼容性说明：该实现不依赖数组最后一维的固定长度，只要Ii0和Iv0前序维度形状匹配即可正常运行。

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