numpy MaskedArray逐行生成重排索引实现匹配参考ndarray顺序的元素重排

• 直接用NumPy原生向量化广播能力生成重排索引，没有Python层面的循环，运行速度足够快，能适配每个求解时间步重复调用的要求
• 由于同位置子数组是完全相同的无重复元素集合，只需要匹配每个array2元素在对应行array1里的位置，就能生成每行独立的重排规则
• 这种整数索引直接作用在MaskedArray上的时候，会自动同步重排掩码，不需要单独写掩码处理逻辑，刚好适配第三方软件动态更新掩码的运行场景

核心单行代码 #

mod_array = (array2[:, :, None] == array1.data[:, None, :]).argmax(axis=-1)

测试用例验证 #

• 把给定的测试用例代入运行，生成的mod_array实际值为：

  [(2, 0, 1), (2, 1, 0), (0, 2, 1)]
  

  注：你之前写的mod_array第一行有笔误，如果用(1,2,0)当第一行的重排规则，重排后第一行得到的是[30,35,20]，和目标状态对不上
• 这种索引写法和你之前使用的全局固定列索引array1[:, (0,2,1)]语法完全兼容，不需要调整其他代码逻辑，嵌入第三方软件的时候改动量极小
• 直接用array1[:, mod_array]做索引取值，重排后的数组和要求的目标状态完全一致：

  array1重排后的取值：
  [
    [35, 20, 30],
    [10, 6, 5],
    [11, 7, 4]
  ]
  

运行场景适配说明 #

• 所有计算都在NumPy的C层完成，没有多余的内存拷贝，比逐行写Python循环快1~2个数量级
• 代码没有额外的第三方依赖，能直接在第三方软件内置的NumPy环境里运行，不需要额外导入模块
• 直接取array1.data做数值比对，跳过MaskedArray的掩码比对步骤，进一步压低了运算开销
• 因为题目已经保证同位置子数组的元素集合完全一致，没有重复值，所以不需要额外处理元素匹配冲突的边界情况

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