寻找一组索引，将一个NumPy ndarray的行映射到另一个数组。

Question

我有两个结构化的2D numpy 数组，这两个数组在原则上是相等的，也就是说

A = numpy.array([[a1,b1,c1],
                 [a2,b2,c2],
                 [a3,b3,c3],
                 [a4,b4,c4]]) 

B = numpy.array([[a2,b2,c2],
                 [a4,b4,c4],
                 [a3,b3,c3],
                 [a1,b1,c1]])

并不是指

numpy.array_equal(A,B) # False
numpy.array_equiv(A,B) # False
numpy.equal(A,B) # ndarray of True and False

但从某种意义上说，一个数组 (A) 是“原始的”，而在另一个数组 (B) 中，数据沿着一个轴（可以是行或列）进行了混洗。

有什么高效的方法可以对B进行排序/洗牌，使其匹配或等于A，或者将A排序以变成B？实际上，平等检查并不重要，只要两个数组被混洗以相互匹配即可。A和因此B具有唯一的行。

我尝试使用view方法对两个数组进行排序，就像这样：

def sort2d(A):
    A_view = np.ascontiguousarray(A).view(np.dtype((np.void,
             A.dtype.itemsize * A.shape[1])))
    A_view.sort()
    return A_view.view(A.dtype).reshape(-1,A.shape[1])   

但这种方法显然在这里行不通。需要对非常大的数组执行此操作，因此性能和可扩展性至关重要。

- romeric

---

B[:] = A 有什么问题吗？ - user2357112

我不能这样做，因为A和B的行分别映射到数组C和D，而A和B的（行）顺序决定了C和D中的值，所以我不能搞乱顺序。 - romeric

Answer 1

根据您的示例，似乎您同时随机排列了所有列，使得有一个行索引向量将A→B进行映射。这里有一个玩具示例：

A = np.random.permutation(12).reshape(4, 3)
idx = np.random.permutation(4)
B = A[idx]

print(repr(A))
# array([[ 7, 11,  6],
#        [ 4, 10,  8],
#        [ 9,  2,  0],
#        [ 1,  3,  5]])

print(repr(B))
# array([[ 1,  3,  5],
#        [ 4, 10,  8],
#        [ 7, 11,  6],
#        [ 9,  2,  0]])

我们想要恢复一组索引idx，使得A[idx] == B。当且仅当A和B不包含重复的行时，这将是一个唯一的映射。

---

一种高效的方法是找到能够在字典排序后排列A中行的索引，然后找到B中每一行在排序后的A中的位置。一个有用的技巧是使用np.void数据类型将A和B视为1D数组，以将每一行视为单个元素：

rowtype = np.dtype((np.void, A.dtype.itemsize * A.size / A.shape[0]))
# A and B must be C-contiguous, might need to force a copy here
a = np.ascontiguousarray(A).view(rowtype).ravel()
b = np.ascontiguousarray(B).view(rowtype).ravel()

a_to_as = np.argsort(a)     # indices that sort the rows of A in lexical order

现在我们可以使用 np.searchsorted 来执行二分查找，以确定每一行在排序后的 A 中的位置：

# using the `sorter=` argument rather than `a[a_to_as]` avoids making a copy of `a`
as_to_b = a.searchsorted(b, sorter=a_to_as)

从A→B的映射可以表示为A→A_s→B的复合。

a_to_b = a_to_as.take(as_to_b)
print(np.all(A[a_to_b] == B))
# True

如果A和B中没有重复的行，那么可以使用以下方法获取从B→A的反向映射。

b_to_a = np.argsort(a_to_b)
print(np.all(B[b_to_a] == A))
# True

作为单个函数：

def find_row_mapping(A, B):
    """
    Given A and B, where B is a copy of A permuted over the first dimension, find
    a set of indices idx such that A[idx] == B.
    This is a unique mapping if and only if there are no repeated rows in A and B.

    Arguments:
        A, B:   n-dimensional arrays with same shape and dtype
    Returns:
        idx:    vector of indices into the rows of A
    """

    if not (A.shape == B.shape):
        raise ValueError('A and B must have the same shape')
    if not (A.dtype == B.dtype):
        raise TypeError('A and B must have the same dtype')

    rowtype = np.dtype((np.void, A.dtype.itemsize * A.size / A.shape[0]))
    a = np.ascontiguousarray(A).view(rowtype).ravel()
    b = np.ascontiguousarray(B).view(rowtype).ravel()
    a_to_as = np.argsort(a)
    as_to_b = a.searchsorted(b, sorter=a_to_as)

    return a_to_as.take(as_to_b)

基准测试：

In [1]: gen = np.random.RandomState(0)
In [2]: %%timeit A = gen.rand(1000000, 100); B = A.copy(); gen.shuffle(B)
....: find_row_mapping(A, B)
1 loop, best of 3: 2.76 s per loop

---

*最耗费时间的步骤是对行进行快速排序，平均时间复杂度为O(n log n)。我不确定是否有更好的方法。

- ali_m

---

我也是这么想的，但你确定这种方法适用于任何一对数组吗？这是在我的笔记本上运行你的代码的两个连续结果：one, two。顺便说一下，我认为你在倒数第二行的意思是 idx 而不是 perm。 - romeric

抱歉，我搞混了正向和反向映射。现在应该可以工作了。 - ali_m

没问题。我尝试了一种替代方法，使用Jaime的答案，通过找到两个数组“vstack”在一起的唯一行（如下所示），虽然它似乎不像你的解决方案那样优雅。 - romeric

Answer 2

由于两个数组中的任何一个都可以洗牌以匹配另一个，因此没有人阻止我们重新排列两个数组。使用 Jaime's Answer，我们可以 vstack 这两个数组并找到唯一的行。然后，unique返回的逆指数实质上是所需的映射（因为数组不包含重复行）。

首先让我们定义一个 unique2d 函数以方便使用：

def unique2d(arr,consider_sort=False,return_index=False,return_inverse=False): 
    """Get unique values along an axis for 2D arrays.

        input:
            arr:
                2D array
            consider_sort:
                Does permutation of the values within the axis matter? 
                Two rows can contain the same values but with 
                different arrangements. If consider_sort 
                is True then those rows would be considered equal
            return_index:
                Similar to numpy unique
            return_inverse:
                Similar to numpy unique
        returns:
            2D array of unique rows
            If return_index is True also returns indices
            If return_inverse is True also returns the inverse array 
            """

    if consider_sort is True:
        a = np.sort(arr,axis=1)
    else:
        a = arr
    b = np.ascontiguousarray(a).view(np.dtype((np.void, 
            a.dtype.itemsize * a.shape[1])))

    if return_inverse is False:
        _, idx = np.unique(b, return_index=True)
    else:
        _, idx, inv = np.unique(b, return_index=True, return_inverse=True)

    if return_index == False and return_inverse == False:
        return arr[idx]
    elif return_index == True and return_inverse == False:
        return arr[idx], idx
    elif return_index == False and return_inverse == True:
        return arr[idx], inv
    else:
        return arr[idx], idx, inv

我们现在可以将我们的映射定义如下：

def row_mapper(a,b,consider_sort=False):
    """Given two 2D numpy arrays returns mappers idx_a and idx_b 
        such that a[idx_a] = b[idx_b] """

    assert a.dtype == b.dtype
    assert a.shape == b.shape

    c = np.concatenate((a,b))
    _, inv = unique2d(c, consider_sort=consider_sort, return_inverse=True)
    mapper_a = inv[:b.shape[0]]
    mapper_b = inv[b.shape[0]:]

    return np.argsort(mapper_a), np.argsort(mapper_b) 

验证:

n = 100000
A = np.arange(n).reshape(n//4,4)
B = A[::-1,:]

idx_a, idx_b  = row_mapper(A,B)
print np.all(A[idx_a]==B[idx_b])
# True

基准测试： 针对 @ali_m 的解决方案进行基准测试

%timeit find_row_mapping(A,B) # ali_m's solution
%timeit row_mapper(A,B) # current solution

# n = 100
100000 loops, best of 3: 12.2 µs per loop
10000 loops, best of 3: 47.3 µs per loop

# n = 1000
10000 loops, best of 3: 49.1 µs per loop
10000 loops, best of 3: 148 µs per loop

# n = 10000
1000 loops, best of 3: 548 µs per loop
1000 loops, best of 3: 1.6 ms per loop

# n = 100000
100 loops, best of 3: 6.96 ms per loop
100 loops, best of 3: 19.3 ms per loop

# n = 1000000
10 loops, best of 3: 160 ms per loop
1 loops, best of 3: 372 ms per loop

# n = 10000000
1 loops, best of 3: 2.54 s per loop
1 loops, best of 3: 5.92 s per loop

虽然当前的解决方案也许还有改进的空间，但是与ali_m的解决方案相比，当前的方案要慢2-3倍，而且可能会更加混乱，因为两个数组都需要映射。只是想提供一个备选方案。

- romeric

---

np.unique 内部使用 argsort。性能差异可能在于您正在对两个数组进行排序而不仅仅是一个，以及一些额外的输入复制。 - ali_m