Question

假设给定两个具有相同行数的2D numpy数组a和b。此外，假设我们知道i和a的每一行b最多具有一个共同的元素，尽管该元素可能出现多次。我们如何才能尽可能高效地找到这个元素？

一个例子：

import numpy as np

a = np.array([[1, 2, 3],
              [2, 5, 2],
              [5, 4, 4],
              [2, 1, 3]])

b = np.array([[4, 5],
              [3, 2],
              [1, 5],
              [0, 5]])

desiredResult = np.array([[np.nan],
                          [2],
                          [5],
                          [np.nan]])

通过沿第一个轴应用intersect1d可以很容易地提出一个有力的实现：

from intertools import starmap

desiredResult = np.array(list(starmap(np.intersect1d, zip(a, b))))

显然，使用python的内置set操作甚至更快。将结果转换为所需形式很容易。

但是，我需要一个尽可能高效的实现。因此，我不喜欢starmap，因为我认为它需要对每一行进行python调用。我想要一个纯粹的向量化选项，并且很高兴，即使这甚至利用了我们的额外知识，即每行最多有一个公共值。

是否有人对我如何加快任务速度和更优雅地实现解决方案有想法？我可以使用C代码或cython，但编码工作不应太多。

Answer 1

方法1

这是基于searchsorted2d的矢量化图标-

# Sort each row of a and b in-place
a.sort(1)
b.sort(1)

# Use 2D searchsorted row-wise between a and b
idx = searchsorted2d(a,b)

# "Clip-out" out of bounds indices
idx[idx==a.shape[1]] = 0

# Get mask of valid ones i.e. matches
mask = np.take_along_axis(a,idx,axis=1)==b

# Use argmax to get first match as we know there's at most one match
match_val = np.take_along_axis(b,mask.argmax(1)[:,None],axis=1)

# Finally use np.where to choose between valid match 
# (decided by any one True in each row of mask)
out = np.where(mask.any(1)[:,None],match_val,np.nan)

方法2

基于Numba的内存效率-

from numba import njit

@njit(parallel=True)
def numba_f1(a,b,out):
    n,a_ncols = a.shape
    b_ncols = b.shape[1]
    for i in range(n):
        for j in range(a_ncols):
            for k in range(b_ncols):
                m = a[i,j]==b[i,k]
                if m:
                    break
            if m:
                out[i] = a[i,j]
                break
    return out

def find_first_common_elem_per_row(a,b):
    out = np.full(len(a),np.nan)
    numba_f1(a,b,out)
    return out

方法3

这是另一个基于堆栈和排序的矢量化-

r = np.arange(len(a))
ab = np.hstack((a,b))
idx = ab.argsort(1)
ab_s = ab[r[:,None],idx]
m = ab_s[:,:-1] == ab_s[:,1:]
m2 = (idx[:,1:]*m)>=a.shape[1]
m3 = m & m2
out = np.where(m3.any(1),b[r,idx[r,m3.argmax(1)+1]-a.shape[1]],np.nan)

方法4

作为一种优雅的方法，我们可以使用broadcasting来节省资源-

m = (a[:,None]==b[:,:,None]).any(2)
out = np.where(m.any(1),b[np.arange(len(a)),m.argmax(1)],np.nan)

Answer 2

做一些研究，我发现检查两个列表是否不相交是在 O（n + m）中进行的，从而 n 和 m 是列表的长度（请参见here）。这个想法是元素的插入和查找对于哈希映射在恒定的时间内运行。因此，将第一个列表中的所有元素插入哈希图中需要进行 O（n）操作，而检查第二个列表中的每个元素是否已经在哈希图中都需要 O（m ）操作。因此，在 O（n log（n）+ m log（m））中运行的基于排序的解决方案并不是渐近最优的。

尽管@Divakar的解决方案在许多用例中都是高效的，但如果第二维很大，它们的效率就会降低。然后，基于哈希图的解决方案更适合。我已经在cython中实现了它：

import numpy as np
cimport numpy as np
import cython
from libc.math cimport NAN
from libcpp.unordered_map cimport unordered_map
np.import_array()

@cython.boundscheck(False)
@cython.wraparound(False)
def get_common_element2d(np.ndarray[double, ndim=2] arr1, 
                         np.ndarray[double, ndim=2] arr2):

    cdef np.ndarray[double, ndim=1] result = np.empty(arr1.shape[0])
    cdef int dim1 = arr1.shape[1]
    cdef int dim2 = arr2.shape[1]
    cdef int i, j
    cdef unordered_map[double, int] tmpset = unordered_map[double, int]()

    for i in range(arr1.shape[0]):
        for j in range(dim1):
            # insert arr1[i, j] as key without assigned value
            tmpset[arr1[i, j]]
        for j in range(dim2):
            # check whether arr2[i, j] is in tmpset
            if tmpset.count(arr2[i,j]):
                result[i] = arr2[i,j]
                break
        else:
            result[i] = NAN
        tmpset.clear()

    return result

我创建了如下测试用例：

import numpy as np
import timeit
from itertools import starmap
from mycythonmodule import get_common_element2d

m, n = 3000, 3000
a = np.random.rand(m, n)
b = np.random.rand(m, n)

for i, row in enumerate(a):
    if np.random.randint(2):
        common = np.random.choice(row, 1)
        b[i][np.random.choice(np.arange(n), np.random.randint(min(n,20)), False)] = common

# we need to copy the arrays on each test run, otherwise they 
# will remain sorted, which would bias the results

%timeit [set(aa).intersection(bb) for aa, bb in zip(a.copy(), b.copy())]
# returns 3.11 s ± 56.8 ms

%timeit list(starmap(np.intersect1d, zip(a.copy(), b.copy)))
# returns 1.83 s ± 55.4

# test sorting method
# divakarsMethod1 is the appraoch #1 in @Divakar's answer
%timeit divakarsMethod1(a.copy(), b.copy())
# returns 1.88 s ± 18 ms

# test hash map method
%timeit get_common_element2d(a.copy(), b.copy())
# returns 1.46 s ± 22.6 ms

这些结果似乎表明，幼稚的方法实际上比某些矢量化版本更好。但是，如果考虑到许多行和更少的列（不同的用例），矢量化算法将发挥其优势。在这种情况下，矢量化方法比朴素的方法快5倍以上，并且排序方法被证明是最好的。

结论：我将使用基于HashMap的cython版本，因为它是这两种用例中最有效的变体之一。如果必须先设置cython，则可以使用基于排序的方法。

Answer 3

不确定这是否更快，但是我们可以在此处尝试一些操作：

方法1 np.intersect1d，具有列表理解功能

[np.intersect1d(arr[0], arr[1]) for arr in list(zip(a,b))]

# Out
[array([], dtype=int32), array([2]), array([5]), array([], dtype=int32)]

或列出：

[np.intersect1d(arr[0], arr[1]).tolist() for arr in list(zip(a,b))]

# Out
[[], [2], [5], []]

方法2 set，具有列表理解功能：

[list(set(arr[0]) & set(arr[1])) for arr in list(zip(a,b))]

# Out
[[], [2], [5], []]

Numpy：有效地找到按行的公共元素

3 个答案: