如何在python中创建包含/排除数据帧中的行的标志

如何在python中创建一个包含/排除df中的行的标志。示例-

flag=1

将包括df中的所有行，

flag=0

将包括除

col2='a'之外的所有行。

row col1 col2
1    1     a
2    2     a
3    3     b
4    4     c

预期输出-

flag=1

应填充所有行，

flag=0

应仅填充行3,4

---

您可以在字典中创建一组带有字符串的规则，并使用

df.query

方法应用它们

import numpy as np
import pandas as pd
np.random.seed(0)

df = pd.DataFrame({
    "col1": np.random.randint(4, size=10), 
    "col2": np.random.choice(list("abcd"), 10)})

print(df)
   col1 col2
0     0    b
1     3    c
2     1    a
3     0    d
4     3    c
5     3    a
6     3    a
7     3    a
8     1    c
9     3    b

创建标志->规则的字典以及应用这些规则的函数：

def flag_filter(df, flag, rule_dict):
    rule = rule_dict[flag]
    if rule is None:
        return df
    return df.query(rule)

rule_dict = {
    0: None,                             # flag 0 returns the dataframe unfiltered
    1: "col2 != 'a'",                    # flag 1 = where column2 is not equal to "a"
    2: "col1 == 3 & col2 == ['b', 'c']"  # flag 2 = where column1 equals 3 AND column2 equals either "b" or "c"
}

标志0

subset = flag_filter(df, 0, rule_dict)
print(subset)

   col1 col2
0     0    b
1     3    c
2     1    a
3     0    d
4     3    c
5     3    a
6     3    a
7     3    a
8     1    c
9     3    b

标志1

subset = flag_filter(df, 1, rule_dict)
print(subset)

   col1 col2
0     0    b
1     3    c
3     0    d
4     3    c
8     1    c
9     3    b

旗帜2

subset = flag_filter(df, 2, rule_dict)
print(subset)

   col1 col2
1     3    c
4     3    c
9     3    b