获取数组中匹配元素的索引,考虑重复

这是一个O(n)-ish解决方案(ish,因为如果重复很长,它显然不可能是O(n))。实际上,根据输入长度的不同,可以通过牺牲O(n)和替换 heapq.merge 有一个马厩 np.argsort . 当N=10^6时,它需要大约1秒的时间。

代码:

import numpy as np

USE_HEAPQ = True

def sqlwhere(a, b):
    asw = np.r_[0, 1 + np.flatnonzero(a[:-1]!=a[1:]), len(a)]
    bsw = np.r_[0, 1 + np.flatnonzero(b[:-1]!=b[1:]), len(b)]
    al, bl = np.diff(asw), np.diff(bsw)
    na, nb = len(al), len(bl)
    abunq = np.r_[a[asw[:-1]], b[bsw[:-1]]]
    if USE_HEAPQ:
        from heapq import merge
        m = np.fromiter(merge(range(na), range(na, na+nb), key=abunq.__getitem__), int, na+nb)
    else:
        m = np.argsort(abunq, kind='mergesort')
    mv = abunq[m]
    midx = np.flatnonzero(mv[:-1]==mv[1:])
    ai, bi = m[midx], m[midx+1] - na
    aic = np.r_[0, np.cumsum(al[ai])]
    a_idx = np.ones((aic[-1],), dtype=int)
    a_idx[aic[:-1]] = asw[ai]
    a_idx[aic[1:-1]] -= asw[ai[:-1]] + al[ai[:-1]] - 1
    a_idx = np.repeat(np.cumsum(a_idx), np.repeat(bl[bi], al[ai]))
    bi = np.repeat(bi, al[ai])
    bic = np.r_[0, np.cumsum(bl[bi])]
    b_idx = np.ones((bic[-1],), dtype=int)
    b_idx[bic[:-1]] = bsw[bi]
    b_idx[bic[1:-1]] -= bsw[bi[:-1]] + bl[bi[:-1]] - 1
    b_idx = np.cumsum(b_idx)
    return a_idx, b_idx

def f_D(a, b):
    return np.where(np.equal.outer(a,b))

def mock_data(n):
    return np.cumsum(np.random.randint(0, 3, (2, n)), axis=1)


a = np.array([3, 4, 4, 7, 9, 9], dtype=np.uint8)
b = np.array([1, 4, 7, 9, 9], dtype=np.uint8)

# check correct
a, b = mock_data(1000)
ai0, bi0 = f_D(a, b)
ai1, bi1 = sqlwhere(a, b)
print(np.all(ai0 == ai1), np.all(bi0 == bi1))

# check fast
a, b = mock_data(1000000)
sqlwhere(a, b)

方法#1:基于Broadacasting的方法

使用 outer 两个数组之间的相等性比较以利用矢量化 broadcasting 然后得到行、列索引,这是两个数组对应的匹配索引所需要的-

a_idx, b_idx = np.where(a[:,None]==b)
a_idx, b_idx = np.where(np.equal.outer(a,b))

我们还可以使用 np.nonzero 代替 np.where .

方法#2:具体案例解决方案

由于没有重复和排序的输入数组,我们可以使用 np.searchsorted ,就像这样-

idx0 = np.searchsorted(a,b)
idx1 = np.searchsorted(b,a)
idx0[idx0==len(a)] = 0
idx1[idx1==len(b)] = 0

a_idx = idx0[a[idx0] == b]
b_idx = idx1[b[idx1] == a]

稍微修改一下,可能更有效的方法是-

idx0 = np.searchsorted(a,b)
idx0[idx0==len(a)] = 0

a_idx = idx0[a[idx0] == b]
b_idx = np.searchsorted(b,a[a_idx])

方法#3:一般情况

以下是一般情况的解决方案(允许重复)-

def findwhere(a, b):
    c = np.bincount(b, minlength=a.max()+1)[a]
    a_idx1 = np.repeat(np.flatnonzero(c),c[c!=0])
    
    b_idx1 = np.searchsorted(b,a[a_idx1])
    m1 = np.r_[False,a_idx1[1:] == a_idx1[:-1],False]
    idx11 = np.flatnonzero(m1[1:] != m1[:-1])
    id_arr = m1.astype(int)
    id_arr[idx11[1::2]+1] = idx11[::2]-idx11[1::2]
    b_idx1 += id_arr.cumsum()[:-1]
    return a_idx1, b_idx1

计时

使用 mock_data 来自@Paul Panzer的soln,用于设置输入:

In [295]: a, b = mock_data(1000000)

# @Paul Panzer's soln
In [296]: %timeit sqlwhere(a, b) # USE_HEAPQ = False
10 loops, best of 3: 118 ms per loop

# Approach #3 from this post
In [297]: %timeit findwhere(a,b)
10 loops, best of 3: 61.7 ms per loop

将重新排列(uint8数据)转换为 1D 阵列

def convert_recarrays_to_1Darrs(a, b):
    a2D = a.view('u1').reshape(-1,2)
    b2D = b.view('u1').reshape(-1,2)
    s = max(a2D[:,0].max(), b2D[:,0].max())+1
    
    a1D = s*a2D[:,1] + a2D[:,0]
    b1D = s*b2D[:,1] + b2D[:,0]
    return a1D, b1D

样本运行-

In [90]: dt = np.dtype([('f1', np.uint8), ('f2', np.uint8)])
    ...: a = np.rec.fromarrays([[3, 4, 4, 7, 9, 9],
    ...:                        [1, 5, 5, 4, 2, 2]], dtype=dt)
    ...: b = np.rec.fromarrays([[1, 4, 7, 9, 9],
    ...:                        [7, 5, 4, 2, 2]], dtype=dt)

In [91]: convert_recarrays_to_1Darrs(a, b)
Out[91]: 
(array([13, 54, 54, 47, 29, 29], dtype=uint8),
 array([71, 54, 47, 29, 29], dtype=uint8))

覆盖的通用版本 rec-arrays

版本#1:

def findwhere_generic_v1(a, b):
    cidx = np.flatnonzero(np.r_[True,b[1:] != b[:-1],True])
    count = np.diff(cidx)
    b_starts = b[cidx[:-1]]
    
    a_starts = np.searchsorted(a,b_starts)
    a_starts[a_starts==len(a)] = 0
    
    valid_mask = (b_starts == a[a_starts])
    count_valid = count[valid_mask]
    
    idx2m0 = np.searchsorted(a,b_starts[valid_mask],'right')    
    idx1m0 = a_starts[valid_mask]
    
    id_arr = np.zeros(len(a)+1, dtype=int)
    id_arr[idx2m0] -= 1
    id_arr[idx1m0] += 1
    
    n = idx2m0 - idx1m0
    r1 = np.flatnonzero(id_arr.cumsum()!=0)
    r2 = np.repeat(count_valid,n)
    a_idx1 = np.repeat(r1, r2)
    
    b_idx1 = np.searchsorted(b,a[a_idx1])
    m1 = np.r_[False,a_idx1[1:] == a_idx1[:-1],False]
    idx11 = np.flatnonzero(m1[1:] != m1[:-1])
    id_arr = m1.astype(int)
    id_arr[idx11[1::2]+1] = idx11[::2]-idx11[1::2]
    b_idx1 += id_arr.cumsum()[:-1]
    return a_idx1, b_idx1

版本#2:

def findwhere_generic_v2(a, b):    
    cidx = np.flatnonzero(np.r_[True,b[1:] != b[:-1],True])
    count = np.diff(cidx)
    b_starts = b[cidx[:-1]]
    
    idxx = np.flatnonzero(np.r_[True,a[1:] != a[:-1],True])
    av = a[idxx[:-1]]
    idxxs = np.searchsorted(av,b_starts)
    idxxs[idxxs==len(av)] = 0
    valid_mask0 = av[idxxs] == b_starts
    
    starts = idxx[idxxs]
    stops = idxx[idxxs+1]
    
    idx1m0 = starts[valid_mask0]
    idx2m0 = stops[valid_mask0]  
    
    count_valid = count[valid_mask0]
    
    id_arr = np.zeros(len(a)+1, dtype=int)
    id_arr[idx2m0] -= 1
    id_arr[idx1m0] += 1
    
    n = idx2m0 - idx1m0
    r1 = np.flatnonzero(id_arr.cumsum()!=0)
    r2 = np.repeat(count_valid,n)
    a_idx1 = np.repeat(r1, r2)
    
    b_idx1 = np.searchsorted(b,a[a_idx1])
    m1 = np.r_[False,a_idx1[1:] == a_idx1[:-1],False]
    idx11 = np.flatnonzero(m1[1:] != m1[:-1])
    id_arr = m1.astype(int)
    id_arr[idx11[1::2]+1] = idx11[::2]-idx11[1::2]
    b_idx1 += id_arr.cumsum()[:-1]
    return a_idx1, b_idx1

纯Python方法

发电机理解

另一种纯python实现,带有生成器和列表理解。与您的代码相比,内存效率可能会更高,但与numpy版本相比可能会更慢。对于排序数组,这将更快。

def pywheregen(a, b):

    l = ((ia,ib) for ia,j in enumerate(a) for ib,k in enumerate(b) if j == k)
    a_idx,b_idx = zip(*l)
    return a_idx,b_idx

考虑数组排序的Python for循环

这里有一个替代版本,它使用简单的python for循环,并考虑到数组被排序,因此它只检查需要检查的对。

def pywhere(a, b):

    l = []
    a.sort()
    b.sort()
    match = 0
    for ia,j in enumerate(a):
        ib = match
        while ib < len(b) and j >= b[ib]:
            if j == b[ib]:
                l.append(((ia,ib)))
                if b[match] < b[ib]:
                    match = ib
            ib += 1

    a_ind,b_ind = zip(*l)

    return a_ind, b_ind

计时

我使用@Paul Panzer的mock\u data()函数比较了计时,并将其与 findwhere() 和 f_D() np。@Divakar的外部进近。 findwhere() 仍然表现最好,但 pywhere() 考虑到它是纯python的,情况并不那么糟糕。 pywheregen() 失败,令人惊讶 f\u D() 需要更长的时间。它们都在N=10^6时失败。由于 heapq 单元

In [2]: a, b = mock_data(10000)
In [10]: %timeit -n 10 findwhere(a,b)                                     
10 loops, best of 3: 1.62 ms per loop

In [11]: %timeit -n 10 pywhere(a,b)                                       
10 loops, best of 3: 20.6 ms per loop

In [12]: %timeit pywheregen(a,b)                                          
1 loop, best of 3: 12.7 s per loop

In [13]: %timeit -n 10 f_D(a,b)                                           
10 loops, best of 3: 476 ms per loop

In [14]: a, b = mock_data(1000000)
In [15]: %timeit -n 10 findwhere(a,b)                                     
10 loops, best of 3: 109 ms per loop

In [16]: %timeit -n 10 pywhere(a,b)                                       
10 loops, best of 3: 2.51 s per loop