Question

假设我有这个数组

array([[100,   1],
       [200,   2],
       [300,   3],
       [400,   4],
       [440,   3]])

我有这个列表或1d数组[100,300]。我希望我的操作输出[1,3]。我怎么能在numpy中做到这一点。

我实际上是在Theano中使用这些numpy数组（一个使用gpu加速计算的机器学习库）。我会有很多行。 Numpy数组允许我将它们无缝地用作Theano中的Tensor对象。但是，如果我必须使用字典，我必须在普通的Python中做到这一点，并且我不确定如果我继续使用大数据，这是否会很好。所以我实际上正在寻找一个numpy操作，索引中的一些技巧或类似的东西。

Answer 1

您可以使用np.in1d：

In [12]: arr
Out[12]: 
array([[100,   1],
       [200,   2],
       [300,   3],
       [400,   4],
       [440,   3]])

In [14]: vals = [100, 300]
In [23]: np.in1d(arr[:,0], vals)
Out[23]: array([ True, False,  True, False, False], dtype=bool)

In [24]: arr[np.in1d(arr[:,0], vals), 1]
Out[24]: array([1, 3])

如果您需要为np.in1d的许多不同值调用vals，则可能需要付费才能准备一个字典，因为 arshajii 建议，因为在准备字典后（a O(n)操作的n = len(arr)操作，查找值将是O(m)操作，其中m = len(vals)。

如果n变得非常大，则dict可能需要太多内存。在这种情况下，您可能需要使用np.in1d。

如果索引（键）值都是ints并且幅度很小，那么可以使用NumPy索引技巧在不使用dict的情况下获得O(m)性能：

In [30]: big = np.full(arr[:,0].max()+1, np.nan)

In [31]: big[arr[:,0]] = arr[:,1]

In [32]: big[vals]
Out[32]: array([ 1.,  3.])

准备big是O(n)操作，但索引big[vals]为O(m)。如果arr[:,0].max()很小且关键值为ints，则使用big的优势在于它比使用dict所需的内存更少。

In [33]: %timeit arr[np.in1d(arr[:,0], vals), 1]
10000 loops, best of 3: 21.5 µs per loop

In [34]: %timeit big[vals]
1000000 loops, best of 3: 1.23 µs per loop

与arshajii的解决方案比较：

In [38]: d = dict(arr)
In [40]: %timeit [d[k] for k in vals]
1000000 loops, best of 3: 447 ns per loop

因此，最佳使用方法取决于arr和vals的大小，执行此操作的次数，内存量以及密钥是否小{{} 1}}。您需要对与您的用例相关的数据进行基准测试，以做出正确的决定。

Answer 2

我只是将您的数组转换为字典：

>>> a = array([[100,   1],
...            [200,   2],
...            [300,   3],
...            [400,   4],
...            [440,   3]])
>>> 
>>> keys = [100, 300]
>>> 
>>> d = dict(a)
>>> 
>>> [d[k] for k in keys]
[1, 3]

Answer 3

如果您确定要搜索的所有值实际存在于搜索数组中，您还可以使用np.searchsorted。对于大型阵列，与其他建议相比似乎更快。

s = np.sort(A[:,0])
A[np.searchsorted(s, values), 1]

如果要搜索的数组已经排序，您可以省略排序过程，操作会更快。

如何在第一个元素位于不同列表的二维数组中提取这些行？

3 个答案: