根据另一个不是索引数组的数组中的值,从numpy数组中选择元素

5
假设我有以下两个数组:
a = array([(1, 'L', 74.423088306605), (5, 'H', 128.05441039929008),
       (2, 'L', 68.0581377353869), (0, 'H', 88.15726964130869), 
       (4, 'L', 97.4501582588212), (3, 'H', 92.98550136344437),
       (7, 'L', 87.75945631669309), (6, 'L', 90.43196739694255),
       (8, 'H', 111.13662092749307), (15, 'H', 91.44444608631304),
       (10, 'L', 85.43615908319185), (11, 'L', 78.11685661303494),
       (13, 'H', 108.2841293816308), (17, 'L', 74.43917911042259),
       (14, 'H', 64.41057325770373), (9, 'L', 27.407214746467943),
       (16, 'H', 81.50506434964355), (12, 'H', 97.79700070323196),
       (19, 'L', 51.139258140713025), (18, 'H', 118.34835768605957)], 
      dtype=[('id', '<i4'), ('name', 'S1'), ('value', '<f8')])

b = array([ 0,  3,  5,  8, 12, 13, 14, 15, 16, 18], dtype=int32)

我希望从b中给出元素a的id进行选择。也就是说,b不是一个索引数组,它包含观测值的ids。在numpy中如何实现这一点?谢谢您的帮助。
3个回答

6

使用这个,你应该能够得到你想要的东西。

indeces = [i for i,id in enumerate(a['id']) if id in b]
suba = a[indeces]
print(suba)
>>>array([(5, 'H', 128.05441039929008), (0, 'H', 88.15726964130869),
   (3, 'H', 92.98550136344437), (8, 'H', 111.13662092749307),
   (15, 'H', 91.44444608631304), (13, 'H', 108.2841293816308),
   (14, 'H', 64.41057325770373), (16, 'H', 81.50506434964355),
   (12, 'H', 97.79700070323196), (18, 'H', 118.34835768605957)], 
  dtype=[('id', '<i4'), ('name', '|S1'), ('value', '<f8')])

谢谢!这看起来不错。如果我一段时间内没有看到更好的答案,我会接受这个。 - Curious2learn

5
以下方法对于你的示例数组比Francesco的方法快几倍:
```html

以下方法对于你的示例数组比Francesco的方法快几倍:

```
In [7]: a[np.argmax(a['id'][None, :] == b[:, None], axis=1)]
Out[7]: 
array([(0, 'H', 88.15726964130869), (3, 'H', 92.98550136344437),
       (5, 'H', 128.05441039929008), (8, 'H', 111.13662092749307),
       (12, 'H', 97.79700070323196), (13, 'H', 108.2841293816308),
       (14, 'H', 64.41057325770373), (15, 'H', 91.44444608631304),
       (16, 'H', 81.50506434964355), (18, 'H', 118.34835768605957)], 
      dtype=[('id', '<i4'), ('name', '|S1'), ('value', '<f8')])

In [8]: %timeit a[np.argmax(a['id'][None, :] == b[:, None], axis=1)]
100000 loops, best of 3: 11.6 us per loop

In [9]: %timeit indices = [i for i,id in enumerate(a['id']) if id in b]; a[indices]
10000 loops, best of 3: 66.9 us per loop

为了理解它是如何工作的,请看这个例子:
In [10]: a['id'][None, :] == b[:, None]
Out[10]: 
array([[False, False, False,  True, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False,
        False, False],
    ... # several rows removed 
    [False, False, False, False, False, False, False, False, False,
        False, False, False, False, False, False, False, False, False,
        False,  True]], dtype=bool)

这是一个数组,它有与b中元素数量相同的行和与a中元素数量相同的列。然后np.argmax在每一行中找到第一个True的位置,这就是a['id']中对应元素的第一次出现的索引。
以上所示,对于小型数组而言,这比Python在性能方面更高效。但是,如果ab其中之一过大,那么布尔值中间数组的大小可能会影响性能。另外,np.argmax必须搜索整行,它永远不会提前跳出循环,如果a太长,这并不是一件好事。我在回答这个问题时进行了一些计时,使用了类似的方法,在处理适度大的数组时仍然是可行的。
Francesco的方法无疑比较简单易懂,而且对于您样本大小的数组而言,性能差异可以忽略不计,我必须承认。但它不会让你感觉像这样...

哇,这太棒了,虽然我不能说我理解 [None,:] 背后的逻辑。只是好奇:你有没有想过你的方法的缩放性?天真地说,我认为我的方法大致与 a 和 b 的大小成线性关系(如果 if id in b 是懒加载,则扩展效果更好)。 - Francesco Montesano
@FrancescoMontesano 这正是问题所在,我认为这是O(n**2),而你的方法更好,尽管根据这个你可能需要将b转换为一个set才能实现。因此,最终你的方法将是最快的,但对于一些较小范围内的数据,Python的速度或者Numpy/C的速度都太快了,以至于这并不重要。 - Jaime
@FrancescoMontesano [None, :] 相当于 .reshape(1, -1),它将一个一维数组转换为列向量。因此,当它比较列向量和行向量时,会将它们广播到完整的矩形形状。 - Jaime
抓住了!我(认为我)不知道这个广播。感谢解释。 - Francesco Montesano

0
sorted = numpy.sort(a)
sorted[b]
 array([(0, 'H', 88.15726964130869), (3, 'H', 92.98550136344437),
   (5, 'H', 128.05441039929008), (8, 'H', 111.13662092749307),
   (12, 'H', 97.79700070323196), (13, 'H', 108.2841293816308),
   (14, 'H', 64.41057325770373), (15, 'H', 91.44444608631304),
   (16, 'H', 81.50506434964355), (18, 'H', 118.34835768605957)], 
  dtype=[('id', '<i4'), ('name', '|S1'), ('value', '<f8')])

只要数组中的行数与ID数量相同即可。

我不想依赖于排序。我希望它能够适应任何顺序。 - Curious2learn
我认为第一行应该是sorted=numpy.argsort(a),这样它就能够适应不同的排序方式。 - DanB

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接