按ID分组实现res列特定加权求和运算的方法咨询

现有如下三列数据：

ID     t      res
     1      1     -1.5
     1      2     -1.5
     1      3      0.5
     1      4      0.5
     2      1     -0.5
     2      2     -0.5
     2      3     -2.0
     2      4     -1.5
     2      5      1.5

需要按ID对res列执行特定运算，最终结果为 -0.167，运算逻辑拆解：

• 对每个ID组：
  1. 遍历组内每个res值，用当前值乘以它之后所有res值的总和
  2. 将所有乘积相加得到组内总乘积和
  3. 把组内总乘积和除以「组内行数-1」
• 将所有组的计算结果相加，得到最终值

Python Pandas 实现 #

基础循环版 #

import pandas as pd

# 构造数据
data = {
    'ID': [1,1,1,1,2,2,2,2,2],
    't': [1,2,3,4,1,2,3,4,5],
    'res': [-1.5,-1.5,0.5,0.5,-0.5,-0.5,-2.0,-1.5,1.5]
}
df = pd.DataFrame(data)

def calculate_group(group):
    res_list = group['res'].tolist()
    group_total = 0
    n = len(res_list)
    for i in range(n):
        # 计算当前值之后的所有res的和
        sum_rest = sum(res_list[i+1:])
        group_total += res_list[i] * sum_rest
    # 除以组内行数-1
    return group_total / (n - 1)

# 按ID分组计算，再求和得到最终结果
final_result = df.groupby('ID').apply(calculate_group).sum()
print(round(final_result, 3))  # 输出 -0.167

向量化优化版（更高效） #

import pandas as pd
import numpy as np

df = pd.DataFrame(data)

def calculate_group_vectorized(group):
    res_arr = group['res'].values
    n = len(res_arr)
    # 计算后缀和（从当前位置到末尾的和，取i+1到末尾的部分）
    suffix_sums = np.cumsum(res_arr[::-1])[::-1][1:]
    # 计算乘积和
    total = np.sum(res_arr[:-1] * suffix_sums)
    return total / (n - 1)

final_result = df.groupby('ID').apply(calculate_group_vectorized).sum()
print(round(final_result, 3))  # 输出 -0.167

R 语言实现（dplyr） #

基础循环版 #

library(dplyr)

# 构造数据
df <- data.frame(
    ID = c(1,1,1,1,2,2,2,2,2),
    t = c(1,2,3,4,1,2,3,4,5),
    res = c(-1.5,-1.5,0.5,0.5,-0.5,-0.5,-2.0,-1.5,1.5)
)

calculate_group <- function(res_vec) {
    group_total <- 0
    n <- length(res_vec)
    for(i in 1:n) {
        sum_rest <- sum(res_vec[(i+1):n])
        group_total <- group_total + res_vec[i] * sum_rest
    }
    group_total / (n - 1)
}

final_result <- df %>%
    group_by(ID) %>%
    summarise(group_val = calculate_group(res)) %>%
    pull(group_val) %>%
    sum()

round(final_result, 3)  # 输出 -0.167

向量化优化版 #

library(dplyr)

calculate_group_vectorized <- function(res_vec) {
    n <- length(res_vec)
    # 计算后缀和
    suffix_sums <- rev(cumsum(rev(res_vec)))[-1]
    total <- sum(res_vec[-n] * suffix_sums)
    total / (n - 1)
}

final_result <- df %>%
    group_by(ID) %>%
    summarise(group_val = calculate_group_vectorized(res)) %>%
    pull(group_val) %>%
    sum()

round(final_result, 3)  # 输出 -0.167

内容的提问来源于stack exchange，提问作者Ahir Bhairav Orai