如何用向量化方法（如np.where）基于另一DataFrame过滤DataFrame？

场景与问题 #

现有两个DataFrame：

df1 #

df2 #

需求：根据df2中对应id的total是否满足条件（如最大值≥25，或存在任意值>25），过滤df1保留符合条件的行。此前使用DataFrame.apply逐行处理，速度极慢，代码如下（注：原代码存在逻辑错误，df_match["id"].max()应为df_match["total"].max()）：

def fn_should_drop(row, check_df):
    df_match = check_df.loc[check_df["id"] == row["id"]]

    max_for_id = df_match["total"].max()  # 修正原代码的字段错误
    max_value = 25
    return max_for_id >= max_value

mask = df1.apply(fn_should_drop, check_df=df2, axis=1)
df_result = df1[mask]

向量化实现方案 #

完全可以用pandas/numpy的向量化操作替代低效的逐行循环，以下是几种高效实现方式：

方式1：基于分组统计的过滤（推荐） #

先对df2按id分组计算统计量，再用符合条件的id过滤df1：

场景A：保留id对应total最大值≥25的行

# 计算每个id的total最大值，筛选出符合条件的id
valid_ids = df2.groupby('id')['total'].max()[lambda x: x >=25].index
# 过滤df1
df_result = df1[df1['id'].isin(valid_ids)]

场景B：保留id对应存在任意total>25的行

# 判断每个id是否存在total>25的记录
valid_ids = df2.groupby('id')['total'].apply(lambda x: x.gt(25).any())[lambda x: x].index
# 过滤df1
df_result = df1[df1['id'].isin(valid_ids)]

方式2：用np.where生成掩码 #

如果需要明确使用np.where，可以先构建id到布尔值的映射，再生成掩码：

# 先获取每个id的有效性映射（以场景A为例）
id_valid_map = df2.groupby('id')['total'].max() >=25
# 用np.where生成掩码
mask = np.where(df1['id'].map(id_valid_map), True, False)
df_result = df1[mask]

性能优势 #

向量化操作基于底层C语言实现，避免了Python层面的逐行循环，在数据量较大时，速度会比apply快10~100倍甚至更多，完全解决原方案的性能问题。

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