如何用列表推导式、np.where/np.select改写Pandas循环计算

原嵌套循环代码（用于计算EQ202207列，数据量较大时耗时30分钟）：

#BOH qty * capacity to get EQ
for i in final_merge2['FLAG']:
    if i =='BOH':
        for j in final_merge2['Plan_ID']:
            if j =='JUL':
                final_merge2['EQ202207'] = (final_merge2['Qty202207']*final_merge2['Capacity'])/1000000

用户尝试的两种错误写法：

• 第一种（np.select语法错误）：

conditions  = [(final_merge2['FLAG'] == 'BOH') & (final_merge2['Plan_ID'] =='JUL')]
choices     = [(final_merge2['Qty202207']*final_merge2['Capacity'])/1000000]
final_merge2["EQ202207"] = np.select(conditions, , default=final_merge2["EQ202207"])

• 第二种（where方法嵌套错误）：

final_merge2["EQ202207"]=((final_merge2['Qty202207']*final_merge2['Capacity'])/1000000).where(final_merge2.FLAG == 'BOH',(.where(final_merge2.Plan_ID=='JUL')))

需求：当final_merge2的FLAG列值为'BOH'且Plan_ID列值为'JUL'时，将Qty202207与Capacity相乘后除以1000000，赋值给EQ202207列；需用列表推导式、np.where或np.select改写原有循环。

方案1：使用np.where（最简洁高效） #

np.where支持向量化操作，是处理单条件赋值的最优选择，符合条件的行执行计算，不符合的可保留原列值或设为NaN：

import numpy as np

# 情况1：EQ202207已存在，保留不符合条件的原有值
final_merge2['EQ202207'] = np.where(
    (final_merge2['FLAG'] == 'BOH') & (final_merge2['Plan_ID'] == 'JUL'),
    (final_merge2['Qty202207'] * final_merge2['Capacity']) / 1000000,
    final_merge2['EQ202207']
)

# 情况2：EQ202207是新增列，不符合条件的设为NaN
final_merge2['EQ202207'] = np.where(
    (final_merge2['FLAG'] == 'BOH') & (final_merge2['Plan_ID'] == 'JUL'),
    (final_merge2['Qty202207'] * final_merge2['Capacity']) / 1000000,
    np.nan
)

方案2：使用np.select #

适合多条件场景，单条件下也可使用，修正了用户之前的语法错误：

import numpy as np

conditions = [(final_merge2['FLAG'] == 'BOH') & (final_merge2['Plan_ID'] == 'JUL')]
choices = [(final_merge2['Qty202207'] * final_merge2['Capacity']) / 1000000]

# 不符合条件时，若列已存在则保留原值，否则设为NaN
final_merge2['EQ202207'] = np.select(conditions, choices, default=final_merge2.get('EQ202207', np.nan))

方案3：使用列表推导式 #

逐行判断逻辑，效率略低于向量化操作，适合小数据量场景：

import numpy as np

final_merge2['EQ202207'] = [
    (qty * cap) / 1000000 if flag == 'BOH' and plan_id == 'JUL' else val
    for qty, cap, flag, plan_id, val in zip(
        final_merge2['Qty202207'],
        final_merge2['Capacity'],
        final_merge2['FLAG'],
        final_merge2['Plan_ID'],
        final_merge2.get('EQ202207', [np.nan]*len(final_merge2))
    )
]

原有错误说明 #

• 第一种np.select写法：np.select的choices参数不能空缺，必须与conditions列表长度一致；
• 第二种where写法：嵌套where缺少调用主体，且未实现“同时满足两个条件”的逻辑判断。

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