检查两个 2D numpy 数组的公共元素，无论是行还是列

Question

给定两个numpy 数组nx3 和mx3，确定行索引（计数器）的有效方法是什么，其中行在两个数组中是共同的。例如，我有以下解决方案，这对于更大的数组来说非常慢

def arrangment(arr1,arr2):
    hits = []
    for i in range(arr2.shape[0]):
        current_row = np.repeat(arr2[i,:][None,:],arr1.shape[0],axis=0)
        x = current_row - arr1
        for j in range(arr1.shape[0]):
            if np.isclose(x[j,0],0.0) and np.isclose(x[j,1],0.0) and np.isclose(x[j,2],0.0):
                hits.append(j)

    return hits

它检查arr1 中是否存在arr2 的行，并返回行匹配的arr1 的行索引。我需要这种安排始终按arr2 的行顺序递增。例如给定

arr1 = np.array([[-1., -1., -1.],
       [ 1., -1., -1.],
       [ 1.,  1., -1.],
       [-1.,  1., -1.],
       [-1., -1.,  1.],
       [ 1., -1.,  1.],
       [ 1.,  1.,  1.],
       [-1.,  1.,  1.]])
arr2 = np.array([[-1.,  1., -1.],
       [ 1.,  1., -1.],
       [ 1.,  1.,  1.],
       [-1.,  1.,  1.]])

函数应该返回：

[3, 2, 6, 7]

Answer 1

快速而肮脏的答案

(arr1[:, None] == arr2).all(-1).argmax(0)

array([3, 2, 6, 7])

---

更好的答案
注意arr2 中的一行与arr1 中的任何内容都不匹配

t = (arr1[:, None] == arr2).all(-1)
np.where(t.any(0), t.argmax(0), np.nan)

array([ 3.,  2.,  6.,  7.])

---

正如@Divakar np.isclose 指出的那样，在比较浮点数时会出现舍入误差

t = np.isclose(arr1[:, None], arr2).all(-1)
np.where(t.any(0), t.argmax(0), np.nan)

Answer 2

我有一个类似的problem in the past，我想出了一个相当优化的解决方案。

首先，您需要将numpy.unique 泛化为多维数组，为了完整起见，我会在此处使用copy

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

现在您只需要连接 (np.vstack) 数组并找到唯一的行。与np.searchsorted 一起反向映射将为您提供所需的索引。因此，让我们编写另一个类似于numpy.in2d 的函数，但用于多维（2D）数组

def in2d_unsorted(arr1, arr2, axis=1, consider_sort=False):
    """Find the elements in arr1 which are also in 
       arr2 and sort them as the appear in arr2"""

    assert arr1.dtype == arr2.dtype

    if axis == 0:
        arr1 = np.copy(arr1.T,order='C')
        arr2 = np.copy(arr2.T,order='C')

    if consider_sort is True:
        sorter_arr1 = np.argsort(arr1)
        arr1 = arr1[np.arange(arr1.shape[0])[:,None],sorter_arr1]
        sorter_arr2 = np.argsort(arr2)
        arr2 = arr2[np.arange(arr2.shape[0])[:,None],sorter_arr2]


    arr = np.vstack((arr1,arr2))
    _, inv = unique2d(arr, return_inverse=True)

    size1 = arr1.shape[0]
    size2 = arr2.shape[0]

    arr3 = inv[:size1]
    arr4 = inv[-size2:]

    # Sort the indices as they appear in arr2
    sorter = np.argsort(arr3)
    idx = sorter[arr3.searchsorted(arr4, sorter=sorter)]

    return idx 

现在你需要做的就是用你的输入参数调用in2d_unsorted

>>> in2d_unsorted(arr1,arr2)
array([ 3,  2,  6,  7])

虽然可能没有完全优化，但这种方法要快得多。让我们将其与 @piRSquareds 解决方案进行基准测试

def indices_piR(arr1,arr2):
    t = np.isclose(arr1[:, None], arr2).all(-1)
    return np.where(t.any(0), t.argmax(0), np.nan)

使用以下数组

n=150
arr1 = np.random.permutation(n).reshape(n//3, 3)
idx = np.random.permutation(n//3)
arr2 = arr1[idx]

In [13]: np.allclose(in2d_unsorted(arr1,arr2),indices_piR(arr1,arr2))
True

In [14]: %timeit indices_piR(arr1,arr2)
10000 loops, best of 3: 181 µs per loop
In [15]: %timeit in2d_unsorted(arr1,arr2)
10000 loops, best of 3: 85.7 µs per loop

现在，为n=1500

In [24]: %timeit indices_piR(arr1,arr2)
100 loops, best of 3: 10.3 ms per loop
In [25]: %timeit in2d_unsorted(arr1,arr2)
1000 loops, best of 3: 403 µs per loop

对于n=15000

In [28]: %timeit indices_piR(A,B)
1 loop, best of 3: 1.02 s per loop
In [29]: %timeit in2d_unsorted(arr1,arr2)
100 loops, best of 3: 4.65 ms per loop

因此，对于大型ish 数组，与@piRSquared 的矢量化解决方案相比，这要快 200 倍。