如何用np.where按条件计算DataFrame行均值？求更优方案

问题背景 #

现有两个pandas DataFrame：

• flag为一维DataFrame，每行包含颜色（red、green、blue）的列表或逗号分隔的颜色字符串
• colors_values存储对应颜色的数值，与flag行数完全匹配

需求：

1. 用np.where根据flag的颜色条件，计算每行对应颜色数值的均值
2. 寻找更高效的替代实现方法

示例数据：

import pandas as pd
import numpy as np

flag = pd.DataFrame({'flag': [ [], ['red'], ['red, green'], ['red, blue'], ['blue'] ]})
colors_values = pd.DataFrame({'red': [1, 1, 1, 1, 1], 'green': [2, 2, 2, 2, 2], 'blue': [4, 4, 4, 4, 4]})

预处理：统一颜色格式 #

首先需要将flag列的内容统一拆分为干净的颜色列表，处理混合格式：

flag['colors'] = flag['flag'].apply(
    lambda x: [c.strip() for c in x[0].split(',')] 
    if isinstance(x, list) and len(x) > 0 and isinstance(x[0], str) 
    else x
)

方法一：使用np.where实现 #

通过逐颜色判断、累加符合条件的数值，再除以有效颜色数量：

# 初始化均值列
flag['mean_value'] = 0.0

# 遍历每个颜色，用np.where筛选并累加数值
for color in colors_values.columns:
    mask = flag['colors'].apply(lambda x: color in x)
    flag['mean_value'] += np.where(mask, colors_values[color], 0)

# 计算有效颜色数，避免除以0
color_count = flag['colors'].apply(len)
flag['mean_value'] = np.where(color_count > 0, flag['mean_value'] / color_count, 0)

# 查看结果
print(flag[['flag', 'mean_value']])

输出结果：

flag  mean_value
0             []         0.0
1          [red]         1.0
2  [red, green]         1.5
3   [red, blue]         2.5
4         [blue]         4.0

方法二：更高效的向量化实现 #

当数据量较大时，逐行apply和循环效率较低，推荐以下向量化方案：

方案A：布尔矩阵匹配（最优性能） #

构建颜色匹配的布尔矩阵，结合数值矩阵批量计算均值：

# 生成每行对应各颜色的选中状态矩阵
color_mask = pd.DataFrame(
    flag['colors'].apply(lambda x: [c in x for c in colors_values.columns]),
    columns=colors_values.columns
)

# 计算均值：布尔矩阵与数值矩阵相乘求和，再除以选中颜色数
mean_values = (color_mask * colors_values).sum(axis=1) / color_mask.sum(axis=1)
# 处理空列表的NaN值
mean_values = mean_values.fillna(0)

flag['mean_value'] = mean_values
print(flag[['flag', 'mean_value']])

方案B：展开分组计算 #

适合需要对颜色做额外处理的场景：

# 展开颜色列表并关联原行索引
expanded = flag.explode('colors').reset_index()
# 合并数值数据
expanded = expanded.merge(colors_values, left_on='index', right_index=True)
# 按原行分组计算均值
mean_values = expanded.groupby('index').apply(
    lambda x: x[x['colors']].mean().mean() if not x['colors'].isna().all() else 0
)

flag['mean_value'] = mean_values.values
print(flag[['flag', 'mean_value']])

性能说明 #

• 小数据量下，np.where方法足够使用；数据量上万行时，方案A的布尔矩阵方法效率最高，完全利用pandas的批量运算能力，避免逐行操作。
• 方案B的灵活性更强，但性能略逊于方案A。

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