python - numpy 数组中元素的顺序

Question

我有一个形状为 (nx3) 的二维数组，比如 arr1。现在考虑第二个数组 arr2，其形状与 arr1 相同并且具有相同的行。但是，行的顺序不同。我想获取 arr2 中每一行的索引，就像它们在 arr1 中一样。我正在寻找最快的 Pythonic 方法来执行此操作，因为 n 的数量级为 10,000。

例如:

arr1 = numpy.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
arr2 = numpy.array([[4, 5, 6], [7, 8, 9], [1, 2, 3]])
ind = [1, 2, 0]

请注意，行元素不必是整数。事实上它们是花车。我找到了使用 numpy.searchsorted 的相关答案，但它们仅适用于一维数组。

Answer 1

如果确保arr2是arr1的排列，则可以使用sort来获取索引:

import numpy as np

n = 100000
a1 = np.random.randint(0, 100, size=(n, 3))
a2 = a1[np.random.permutation(np.arange(n))]
idx1 = np.lexsort(a1.T)
idx2 = np.lexsort(a2.T)
idx = idx2[np.argsort(idx1)]
np.all(a1 == a2[idx])

如果它们没有完全相同的值，您可以在 scipy 中使用 kdTree:

n = 100000

a1 = np.random.uniform(0, 100, size=(n, 3))
a2 = a1[np.random.permutation(np.arange(n))] + np.random.normal(0, 1e-8, size=(n, 3))
from scipy import spatial
tree = spatial.cKDTree(a2)
dist, idx = tree.query(a1)
np.allclose(a1, a2[idx])