Python pandas按ID分组条件过滤并保留原索引的实现问题

解决方案 #

完整可运行的实现代码如下：

import pandas as pd

# 原判断逻辑的向量化实现，适配pandas Series输入
def is_perfect_square(series):
    sqrt_num = series ** 0.5
    return sqrt_num.astype(int) ** 2 == series

# 测试数据
df = {
   'ID':['1','1','1','1','1','2','2','2','2','2','3','3','3','3','3'],
   'Percentage':[7,8,9,10,11,12,13,14,15,16,17,18,19,20,21]}
df = pd.DataFrame(df)

# 核心过滤逻辑
# 按ID分组，判断组内是否存在符合条件的数值，生成每行对应的布尔掩码
mask = df.groupby('ID')['Percentage'].transform(lambda x: is_perfect_square(x).any())
# 用掩码过滤行，保留原始索引
result = df[mask]

运行后result输出与预期完全一致，类型为标准pandas DataFrame，且完全保留原始索引。

补充说明 #

• 整个运算过程为pandas原生向量化操作，没有逐行/逐组调用自定义标量函数，性能远高于普通apply实现，适合处理大数据量场景
• 如果需要复用你原有funct1，可以将判断逻辑替换为transform(lambda x: x.apply(funct1).any())，但该方式不属于纯向量化实现，数据量大时性能会更低，优先推荐上述向量化方案

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