Python中基于DataFrame多条件创建计数列的向量化优化方案咨询

嘿，你的问题我太熟悉了——用iterrows()处理百万行数据简直是性能灾难，完全浪费了Pandas的向量化优势！咱们直接用布尔索引和向量化操作重构代码，速度能提升几十甚至上百倍，而且代码更简洁易维护。

核心思路 #

抛弃逐行迭代的低效方式，利用Pandas的向量化布尔运算批量处理所有行，底层由C引擎执行，避开Python层面的循环开销。同时优化列初始化逻辑，避免不必要的内存操作。

完整实现代码 #

import pandas as pd

# 读取数据
df = pd.read_csv('file.csv')

# 定义需要创建的计数列
tempcols = ['email_cnt', 'event_cnt', 'dm_cnt', 'enc_cnt', 'exp_cnt', 
            'orgsearch_cnt', 'orgsocial_cnt', 'paidsm_cnt', 'paidsearch_cnt', 
            'pd_cnt', 'smrtroom_cnt', 'stry_cnt', 'search_cnt']

# 直接初始化所有计数列为0（比concat空DataFrame高效10倍以上）
df[tempcols] = 0

# 定义通用关注的Status取值
status_vals = ['Accepted', 'Call Completed', 'Commented', 'Declined', 'Liked', 'Responded', 'Shared']

# ---------------------- 处理R1、R2对应的计数列 ----------------------
# email_cnt：Record为R1且Status在指定列表
df['email_cnt'] = ((df['Record'] == 'R1') & df['Status'].isin(['Open', 'Click'])).astype(int)

# event_cnt：Record为R2且Status在指定列表
df['event_cnt'] = ((df['Record'] == 'R2') & df['Status'].isin(['Open', 'Registered'])).astype(int)

# ---------------------- 批量处理R3对应的各Type计数列 ----------------------
# 用字典映射Type到对应计数列和Status列表，避免重复写冗余代码
r3_type_mapping = {
    'Story': ('stry_cnt', status_vals),
    'Search': ('search_cnt', status_vals + ['Downloaded', 'Registered']),
    'Experience': ('exp_cnt', status_vals),
    # 可继续添加其他Type的映射，比如'DM'对应'dm_cnt'、'Encounter'对应'enc_cnt'等
}

# 仅循环Type种类（而非百万行），效率极高
for type_val, (target_col, allowed_status) in r3_type_mapping.items():
    condition = (df['Record'] == 'R3') & (df['Type'] == type_val) & df['Status'].isin(allowed_status)
    df[target_col] = condition.astype(int)

为什么这个方案更快？ #

1. 向量化运算：所有条件判断都是在Pandas底层C引擎中批量执行，避免了Python逐行循环的巨大开销，百万级数据处理时间从分钟级压缩到秒级。
2. 优化列初始化：直接df[tempcols] = 0比concat空DataFrame再填充NaN的逻辑更简洁，内存占用更低。
3. 避免逐行赋值：抛弃iterrows()和df.at[]这类逐行操作，改用批量布尔赋值，彻底发挥Pandas的性能优势。

额外优化建议 #

如果你的Type种类非常多，可以结合pd.get_dummies生成基础计数列后，再通过布尔条件过滤符合Status要求的行，进一步简化代码，但上面的字典映射方式已经足够高效且易读。

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