按指定规则同步排序三个关联数组的实现方案问询

问题需求 #

需要对三个NumPy数组执行排序操作，需满足以下规则：

• 数组array1按递增顺序排列，达到最大值后回到0；
• 数组array2按从小到大的顺序排序；
• 数组array3的每个元素必须始终与前两个数组对应位置的数对保持关联，排序前后映射关系不变。

初始代码示例 #

import numpy as np

array1 = np.array([ 0 ,  0 ,  1 ,  1 ,  2 ])
array2 = np.array([ 1 ,  0 ,  1 ,  0 ,  0 ])
array3 = np.array(['L', 'H', 'O', 'E', 'L'])

期望输出 #

array1 = np.array([ 0 ,  1 ,  2 ,  0 ,  1 ])
array2 = np.array([ 0 ,  0 ,  0 ,  1 ,  1 ])
array3 = np.array(['H', 'E', 'L', 'L', 'O'])

实现方案 #

核心思路是先按array2升序排序，再在array2值相同的分组内按array1升序排序，利用NumPy的lexsort函数实现多键排序，同时保证三个数组的对应关联关系。

完整代码 #

import numpy as np

# 初始数组
array1 = np.array([0, 0, 1, 1, 2])
array2 = np.array([1, 0, 1, 0, 0])
array3 = np.array(['L', 'H', 'O', 'E', 'L'])

# 生成排序索引：优先按array2升序，再按array1升序
sort_indices = np.lexsort((array1, array2))

# 对三个数组应用排序索引
sorted_array1 = array1[sort_indices]
sorted_array2 = array2[sort_indices]
sorted_array3 = array3[sort_indices]

# 打印结果
print("array1 =", sorted_array1)
print("array2 =", sorted_array2)
print("array3 =", sorted_array3)

运行结果 #

array1 = [0 1 2 0 1]
array2 = [0 0 0 1 1]
array3 = ['H' 'E' 'L' 'L' 'O']

原理说明 #

• np.lexsort((array1, array2))函数会按照传入键的逆序进行排序：先以array2作为第一排序键（升序），当array2元素值相同时，再以array1作为第二排序键（升序），最终生成一组排序索引。
• 通过这组索引重新索引三个原数组，即可保证array3的元素与array1、array2的对应数对始终保持映射关系。

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