本文共 890 字，大约阅读时间需要 2 分钟。

pandas :加入有条件的数据框

import pandas as pd# 创建一个数据框df = pd.DataFrame({
   'A': [1, 2, 3, 4, 5], 'B': ['a', 'b', 'c', 'd', 'e']})# 使用query()函数查询满足条件的数据行condition_data = df.query('A > 2')print(condition_data)# 输出：#   A B# 3  4 d# 4  5 e

在这个例子中，我们首先创建了一个包含两列数据的数据框。然后，我们使用query()函数查询满足条件的数据行（A列的值大于2）。最后，我们打印出满足条件的数据行。

如果你想要选择特定的列，你可以直接在.loc[]或iloc[]的参数中指定列名：

# 选择满足条件的行的B列condition_data = df.query('A > 2')['B']print(condition_data)# 输出：# 3    d# 4    e# Name: B, dtype: object

在这个例子中，我们只选择了满足条件的行的B列。

测试用例如下：

df = pd.DataFrame({
   'A': [1, 2, 3, 4, 5], 'B': ['a', 'b', 'c', 'd', 'e']})assert df.query('A > 2').equals(pd.DataFrame({
   'A': [3, 4, 5], 'B': ['c', 'd', 'e']}))

这个测试用例检查了query()函数是否正确地返回了满足条件的行。

如果你的应用场景需要对数据框进行更复杂的条件查询，你可以使用pandas的多个功能来组合不同的查询条件。例如，你可以在一个query()函数中嵌套另一个query()函数，或者在一个query()函数中使用&和|操作符来组合多个条件。