最快符合矩阵

Question

我有两个数组，我想计算一个重合的列表/数组。即所有索引i，j的列表，使得ai == bj。这是我现在的代码：

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

np.array([[i, j] for i in range(a.size) for j in range(b.size) if a[i] == b[j]])

有没有更快的方法(也许是numpy驱动的)来做这件事？

Answer 1

方法#1

一种方法是使用np.in1d -

m_a = np.in1d(a,b)
I = np.flatnonzero(m_a)
J = np.flatnonzero(np.in1d(b, a[m_a]))

样本输入，输出-

In [367]: a
Out[367]: array([1, 2, 3, 4])

In [368]: b
Out[368]: array([3, 5, 6, 4])

In [370]: I
Out[370]: array([2, 3])

In [371]: J
Out[371]: array([0, 3])

方法#2

另一种简单但占用大量内存的方法是使用broadcasting -

I,J = np.nonzero(a[:,None] == b)

方法#3

对于输入数组中没有重复项的情况，我们可以使用np.searchsorted。这里有两个变体--一个用于排序的a，另一个用于通用的a。

变体#1 :排序a的-

idx = np.searchsorted(a, b)
idx[idx==a.size] = 0
mask = a[idx] == b
I = np.searchsorted(a,b[mask])
J = np.flatnonzero(mask)

变体#2 :对于这个通用变体，我们需要使用a的argsort索引-

sidx = a.argsort()
a_sort = a[sidx]
idx = np.searchsorted(a_sort, b)
idx[idx==a.size] = 0
mask = a_sort[idx] == b
I = sidx[np.searchsorted(a_sort,b[mask])]
J = np.flatnonzero(mask)

Answer 2

numpy解决方案可能使用函数numpy.argwhere()，该函数可用于查找符合给定条件的数组的索引。

ax = np.tensordot(a, np.ones(len(a)), axes = 0)
bx = np.tensordot(np.ones(len(b)), b, axes = 0)
np.argwhere(ax - bx == 0)

ax - bx的零元素的索引是那些对应于a和b的相等元素的索引，因为存在常量行rsp。张量积的列“相交”。不过，我不确定这个解决方案是否更快。

Answer 3

已经解决了，但是时间到了

我已经将您的解决方案与python-list方法以及建议的矩阵解决方案here进行了比较：(不过，您需要从该矩阵获取索引)。

import numpy as np
import random
import time

random.seed(12345)
b = [random.randint(0,100000) for i in range(10000)]
a = [random.randint(0,100000) for i in range(10000)]

为了获得更准确的时间(对于大型数据集)，我创建了两个长度为1e5值的(伪)随机整数列表。

#List based approach
start_time = time.time()
c2 = [[i,j]  for i in range(len(a)) for j in range(len(b)) if a[i] == b[j]]
print time.time() - start_time # t = 29.2776758671

b = np.array(a)
a = np.array(b)

#NumPy Array based approach    
start_time = time.time()
c1 = np.array([[i, j] for i in range(a.size) for j in range(b.size) if a[i] == b[j]])
print time.time() - start_time # t = 46.374776125

这并不是一个巨大的时间差，尽管不使用numpy数组所需的时间略短，但对于大向量来说，这仍然意味着相当长的计算时间。

创建一个巧合形式的中间解决方案

#Coincidence Matrix (NumPy) based approach    
start_time = time.time()
c3 = (a[None,:] == b[:,None]).astype(int)
c3s = np.where(c3 == 1)
print time.time() - start_time # t = 0.857568979263

我还对之前发布的另一个解决方案进行了计时，似乎是解决这个问题的最快方法：

c4 = np.nonzero(a[:,None] == b)
# t = 0.399062156677