快速查找集合中的 numpy 向量答案

【问题标题】：Find numpy vectors in a set quickly快速查找集合中的 numpy 向量
【发布时间】：2016-08-30 04:26:16
【问题描述】：

我有一个numpy数组，例如：

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

我也有一套

b = set((1,2),(6,4),(9,9))

我想找到存在于集合 b 中的向量的索引，这里是

[0, 2]

但我使用 for 循环来实现这一点，有没有一种方便的方法可以避免 for 循环？我使用的for循环方法：

record = []
for i in range(a.shape[0]):
    if (a[i, 0], a[i, 1]) in b:
        record.append(i)

【问题讨论】：

作为列表理解...您可以将其解压缩以使其更具可读性... np.vstack([a[np.sum(np.where(a == i, 1, 0 ), axis=1) == 2] for i in b]) ... 但基本上你想找出 b 中的条目在 a 中的位置，当 b 更改为数组时 b = np.array([ (1,2),(6,4)])
该集合中的所有元素是否都保证在a 中的某个位置？
@Divakar 不，它没有
所以，假设输入集中还有一个元素(9,9)。那么输出必须是什么？为什么不在问题的示例中添加这样的案例？
@Divakar 它不会改变答案。我已经更新了这个问题。但我不知道为什么它很重要？

标签： python numpy set vectorization lookup

【解决方案1】：

你可以使用过滤器：

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

In [9]: b = {(1,2),(6,4)}

In [10]: filter(lambda x: tuple(a[x]) in b, range(len(a)))
Out[10]: [0, 2]

【讨论】：

过滤是否比 for 循环快得多？
filter 也将遍历数组。

【解决方案2】：

首先，将集合转换为 NumPy 数组 -

b_arr = np.array(list(b))

然后，基于this post，您将有三种方法。让我们使用第二种方法来提高效率 -

dims = np.maximum(a.max(0),b_arr.max(0)) + 1
a1D = np.ravel_multi_index(a.T,dims)
b1D = np.ravel_multi_index(b_arr.T,dims)    
out = np.flatnonzero(np.in1d(a1D,b1D))

示例运行 -

In [89]: a
Out[89]: 
array([[1, 2],
       [3, 4],
       [6, 4],
       [5, 3],
       [3, 5]])

In [90]: b
Out[90]: {(1, 2), (6, 4), (9, 9)}

In [91]: b_arr = np.array(list(b))

In [92]: dims = np.maximum(a.max(0),b_arr.max(0)) + 1
    ...: a1D = np.ravel_multi_index(a.T,dims)
    ...: b1D = np.ravel_multi_index(b_arr.T,dims)    
    ...: out = np.flatnonzero(np.in1d(a1D,b1D))
    ...: 

In [93]: out
Out[93]: array([0, 2])

【讨论】：

这个方法会运行慢吗？似乎它将使用所有a的元素来比较b的元素。使用set会不会更快，因为是hash方法？
@maple 更新为更高效。这应该快得多。

【解决方案3】：

作为参考，一个简单的列表理解（循环）答案：

In [108]: [i for i,v in enumerate(a) if tuple(v) in b]
Out[108]: [0, 2]

与filter的方法基本一样的速度：

In [111]: timeit [i for i,v in enumerate(a) if tuple(v) in b]
10000 loops, best of 3: 24.5 µs per loop

In [114]: timeit list(filter(lambda x: tuple(a[x]) in b, range(len(a))))
10000 loops, best of 3: 29.7 µs per loop

但这是一个玩具示例，所以时间没有意义。

如果a 还不是一个数组，由于创建数组的开销，这些列表方法将比数组方法更快。

有一些 numpy 集合操作，但它们适用于一维数组。我们可以通过将 2d 数组转换为 1d 结构化来解决这个问题。

In [117]: a.view('i,i')
Out[117]: 
array([[(1, 2)],
       [(3, 4)],
       [(6, 4)],
       [(5, 3)],
       [(3, 5)]], 
      dtype=[('f0', '<i4'), ('f1', '<i4')])
In [119]: np.array(list(b),'i,i')
Out[119]: 
array([(1, 2), (6, 4), (9, 9)], 
      dtype=[('f0', '<i4'), ('f1', '<i4')])

有一个使用 np.void 的版本，但使用这个 'i,i' dtype 更容易记住和玩。

所以这行得通：

In [123]: np.nonzero(np.in1d(a.view('i,i'),np.array(list(b),'i,i')))[0]
Out[123]: array([0, 2], dtype=int32)

但它比迭代慢得多：

In [124]: timeit np.nonzero(np.in1d(a.view('i,i'),np.array(list(b),'i,i')))[0]
10000 loops, best of 3: 153 µs per loop

正如在其他最近的union 问题中所讨论的，np.in1d 使用了几种策略。一种是基于广播和where。其他使用unique、concatenation、sorting和区别。

一个广播解决方案（是的，它很混乱）——但比in1d 更快。

In [150]: timeit np.nonzero((a[:,:,None,None]==np.array(list(b))[:,:]).any(axis=-1).any(axis=-1).all(axis=-1))[0]
10000 loops, best of 3: 52.2 µs per loop

【讨论】：

【解决方案4】：

使用列表理解的单行解决方案：

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

In [63]: b = set(((1,2),(6,4),(9,9)))
In [64]: where([tuple(e) in b for e in a])[0]
Out[64]: array([0, 2])

【讨论】：