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（这些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

这里有一个解决方案，可能比任何 Python 的 for 循环运行得更快。不过，我觉得它的表现可能不如 in1d。只有在你能接受使用一个临时的二维整数数组，大小为 ids.size 乘以 table.IDs.size 的情况下，才可以使用这个方法。在这里，ids 是一个 numpy 数组，里面存放的是 id_list。

result = table[~np.all(table.IDs[None]-ids[None].T, 0)]

Answer 2

你可以这样做：

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索引