Numpy表 - 高级多标准选择

Question

我有一张表是这样的:

IDs    Timestamp     Values

124    300.6          1.23
124    350.1         -2.4
309    300.6          10.3
12     123.4          9.00
18     350.1          2.11
309    350.1          8.3

       ...

我想选择属于一组ID的所有行.我知道我可以做类似的事情

table[table.IDs == 124]

选择一个ID的所有行,我就能做到

table[(table.IDs == 124) | (table.IDs == 309)]

获得两个ID的行.但是想象一下,我有~100,000行,超过1,000个唯一ID(与行索引不同),我想选择匹配一组10个ID的所有行.直觉我想这样做:

# id_list: a list of 10 IDs
table[ table.IDs in id_list ]

但是Python拒绝这种语法.我能想到的唯一方法是做以下事情:

table[ (table.IDs == id_list[0]) |
       (table.IDs == id_list[1]) |
       (table.IDs == id_list[2]) |
       (table.IDs == id_list[3]) |
       (table.IDs == id_list[4]) |
       (table.IDs == id_list[5]) |
       (table.IDs == id_list[6]) |
       (table.IDs == id_list[7]) |
       (table.IDs == id_list[8]) |
       (table.IDs == id_list[9]) ]

这对我来说似乎非常不优雅 - 代码太多而且不同长度的列表没有灵活性.有没有解决我的问题的方法,例如使用列表推导或.any()函数？任何帮助表示赞赏.

Answer 1

你可以这样做:

subset = table[np.array([i in id_list for i in table.IDs])]

如果您有更新版本的numpy,您可以使用该in1d函数使其更紧凑:

subset = table[np.in1d(table.IDs, id_list)]

另请参阅此问题:基于与外部数组交集的numpy recarray索引