如何更简洁高效地获取pandas DataFrame中符合集合匹配条件的列值？

问题背景 #

现有代码可实现需求，但写法繁琐，希望找到更简洁高效的实现方式：

import pandas as pd

data = {
    "A": [1, 2, 3],
    "B": [100, 200, 300]
}

s = {100, 300, 500}

df = pd.DataFrame(data)

for _, row in df.query("B in @s").iterrows():
    print(row["A"])

输出：

1
3

简洁高效的实现方案 #

方案1：用isin直接筛选提取（最推荐）

利用Pandas的isin方法快速匹配，直接提取目标列，避免逐行迭代：

import pandas as pd

data = {
    "A": [1, 2, 3],
    "B": [100, 200, 300]
}

s = {100, 300, 500}

df = pd.DataFrame(data)

# 筛选B列在集合中的行，提取A列
result = df[df["B"].isin(s)]["A"]

# 打印单个值（和原代码输出一致）
for val in result:
    print(val)

# 或者直接输出列表形式，无需循环
print(result.tolist())  # 输出 [1, 3]

方案2：简化query写法

如果偏好query语法，直接在查询后提取A列即可：

result = df.query("B in @s")["A"]
print(result.tolist())

方案3：用loc精准定位

loc方法可以同时指定行条件和列，代码可读性更强：

result = df.loc[df["B"].isin(s), "A"]
print(result.tolist())

效率对比 #

原代码用iterrows逐行迭代，在数据量大时效率极低。上述方案均采用Pandas向量式操作，基于底层优化实现，执行速度远快于逐行循环，同时代码更简洁。

内容的提问来源于stack exchange，提问作者jackal