python - 如何找到重新排序的 numpy 数组的索引？

Question

假设我有一个排序的 numpy 数组:

arr = np.array([0.0, 0.0],
               [0.5, 0.0],
               [1.0, 0.0],
               [0.0, 0.5],
               [0.5, 0.5],
               [1.0, 0.5],
               [0.0, 1.0],
               [0.5, 1.0],
               [1.0, 1.0])

假设我对它做了一个非常重要的操作，这样我就有了一个新数组，它与旧数组相同但顺序不同:

arr2 = np.array([0.5, 0.0],
                [0.0, 0.0],
                [0.0, 0.5],
                [1.0, 0.0],
                [0.5, 0.5],
                [1.0, 0.5],
                [0.0, 1.0],
                [1.0, 1.0],
                [0.5, 1.0])

问题是:如何获取 arr2 的每个元素在 arr 中的位置的索引。换句话说，我想要一个方法，它接受两个数组并返回一个与 arr2 长度相同但具有 arr 元素索引的数组。例如，返回数组的第一个元素将是 arr 中 arr2 的第一个元素的索引。

where_things_are(arr2, arr) 
return : array([1, 0, 3, 2, 4, 5, 6, 8, 7])

numpy 中是否已经存在这样的函数？

编辑:

我试过:

np.array([np.where((arr == x).all(axis=1)) for x in arr2])

它返回我想要的，但我的问题仍然存在:是否有使用 numpy 方法更有效的方法？

编辑 2:

如果 arr2 的长度与原始数组的长度不同(比如我从中删除了一些元素)，它也应该有效。因此，它不是寻找和反转排列，而是寻找元素所在的位置。

Answer 1

关键是反转排列。即使原始数组未排序，下面的代码也能正常工作。如果已排序，则可以使用 find_map_sorted，这显然更快。

更新:为了适应 OP 不断变化的要求，我添加了一个分支来处理丢失的元素。

import numpy as np

def invperm(p):
    q = np.empty_like(p)
    q[p] = np.arange(len(p))
    return q

def find_map(arr1, arr2):
    o1 = np.argsort(arr1)
    o2 = np.argsort(arr2)
    return o2[invperm(o1)]

def find_map_2d(arr1, arr2):
    o1 = np.lexsort(arr1.T)
    o2 = np.lexsort(arr2.T)
    return o2[invperm(o1)]

def find_map_sorted(arr1, arrs=None):
    if arrs is None:
        o1 = np.lexsort(arr1.T)
        return invperm(o1)
    # make unique-able
    rdtype = np.rec.fromrecords(arrs[:1, ::-1]).dtype
    recstack = np.r_[arrs[:,::-1], arr1[:,::-1]].view(rdtype).view(np.recarray)
    uniq, inverse = np.unique(recstack, return_inverse=True)
    return inverse[len(arrs):]

x1 = np.random.permutation(100000)
x2 = np.random.permutation(100000)
print(np.all(x2[find_map(x1, x2)] == x1))

rows = np.random.random((100000, 8))
r1 = rows[x1, :]
r2 = rows[x2, :]
print(np.all(r2[find_map_2d(r1, r2)] == r1))

rs = r1[np.lexsort(r1.T), :]
print(np.all(rs[find_map_sorted(r2), :] == r2))

# lose ten elements
print(np.all(rs[find_map_sorted(r2[:-10], rs), :] == r2[:-10]))