如何找到将一个数组置换为另一个数组的向量化Python函数？

要实现完全向量化的数组置换索引查找，替代带for循环的版本，可利用NumPy的argsort函数完成，无需显式循环。

方法一：双重argsort映射 #

利用argsort的特性，先对x1生成排序索引，再通过x2的排序索引映射得到最终的置换数组：

import numpy as np

def find_permutation_vectorized(x1, x2):
    sorter_x1 = np.argsort(x1)
    # 通过x2的排序索引，映射到x1的原索引
    return sorter_x1[np.argsort(x2)]

测试示例：

x1 = np.array([3, 1, 4, 2])
x2 = np.array([2, 1, 3, 4])
sort = find_permutation_vectorized(x1, x2)
print(sort)
# 输出: array([3, 1, 0, 2])
print(x1[sort])
# 输出: array([2, 1, 3, 4])

方法二：结合searchsorted与argsort #

如果数组元素是有序范围的数值，也可以用searchsorted快速定位元素位置，再结合x1的排序索引得到结果：

def find_permutation_vectorized(x1, x2):
    sorter_x1 = np.argsort(x1)
    sorted_x1 = x1[sorter_x1]
    # 找到x2元素在排序后x1中的位置，再映射回原索引
    pos = np.searchsorted(sorted_x1, x2)
    return sorter_x1[pos]

两种方法均为纯向量化实现，对于大规模数组的处理效率远高于循环版本，且能准确满足需求。

内容的提问来源于stack exchange，提问作者W. K. Black