在Pandas中在for循环中创建布尔索引

Question

我想用布尔索引从一个 pandas 数据框中获取一个子集。

我想测试的条件是这样的：(df[var_0] == value_0) & ... & (df[var_n] == value_n)，其中涉及的变量数量 n 是可以变化的。所以我不能直接写：

df = df[(df[var_0] == value_0) & ... & (df[var_n] == value_n)]

我可以这样做：

for k in range(0,n+1) :
    df = df[df[var_k] == value_k]

（加上一些错误处理，以确保在数据框为空时也能正常工作），但这样似乎不是很高效。有没有人知道怎么用更简洁的 pandas 方式来写这个？

Answer 1

这里的 isin 方法应该能帮到你。

In [7]: df
Out[7]: 
   a  b  c  d  e
0  6  3  1  9  6
1  8  9  5  7  2
2  6  4  7  4  3
3  4  8  0  0  5
4  4  4  2  3  4
5  2  5  9  0  9
6  4  8  2  9  1
7  3  0  8  9  7
8  0  5  9  9  6
9  0  7  8  4  8

[10 rows x 5 columns]

In [8]: vals = {'a': [3], 'b': [0], 'c': [8], 'd': [9], 'e': [7]}

In [9]: df.isin(vals)
Out[9]: 
       a      b      c      d      e
0  False  False  False   True  False
1  False  False  False  False  False
2  False  False  False  False  False
3  False  False  False  False  False
4  False  False  False  False  False
5  False  False  False  False  False
6  False  False  False   True  False
7   True   True   True   True   True
8  False  False  False   True  False
9  False  False   True  False  False

[10 rows x 5 columns]

In [10]: df[df.isin(vals).all(1)]
Out[10]: 
   a  b  c  d  e
7  3  0  8  9  7

[1 rows x 5 columns]

在 vals 字典里的值需要是一个集合，所以我把它们放进了长度为1的列表里。也有可能 query 也能做到这一点。