pandas按country字段groupby后按feature条件计算count值实现问题

实现方案 #

你可以通过pandas的groupby结合自定义聚合逻辑实现需求，以下是两种常用写法： #

写法1：自定义聚合函数（直观易读）

适合数据量不大的场景，逻辑清晰易懂：

import pandas as pd

# 构造原始DataFrame
country = ['US', 'US', 'US', 'UK', 'UK', 'Canada', 'Canada', "Mexico"]
feature =  [2, 2, 2, 1, 1, 2, 2, 1]
ID = [1, 2, 1, 3, 4, 1, 2, 1]  
df = pd.DataFrame(list(zip(country,feature, ID)), columns =['country', 'feature', 'ID'])

# 定义分组计算规则
def calc_count(group):
    # 同一国家的feature值统一，取第一个值判断即可
    if group['feature'].iat[0] == 2:
        # 返回ID的唯一值数量
        return group['ID'].nunique()
    # 返回分组总条目数
    return len(group)

# 分组聚合得到结果
result = df.groupby('country', as_index=False).apply(calc_count)
result.columns = ['country', 'count']
print(result)

写法2：预聚合指标（性能更高）

适合大数据量场景，避免逐行调用自定义函数的性能损耗：

# 先批量聚合每个国家的所有需要的指标
group_agg = df.groupby('country').agg(
    feature_val=('feature', 'first'),
    id_unique=('ID', 'nunique'),
    total_cnt=('country', 'count')
).reset_index()

# 按规则生成count列
group_agg['count'] = group_agg.apply(
    lambda row: row['id_unique'] if row['feature_val'] == 2 else row['total_cnt'],
    axis=1
)

# 提取最终需要的字段
result = group_agg[['country', 'count']].reset_index(drop=True)
print(result)

输出结果

两种写法的输出完全符合预期：

country  count
0      US      2
1      UK      2
2  Canada      2
3  Mexico      1

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