编写R语言函数为data.table新增列：第i行为i+1行起两列乘积之和

问题回顾 #

需要为data.table新增一列z，其中第i行的值为从第i+1行开始的所有x*y的乘积之和。示例数据：

library(data.table)
dt <- data.table(x=1:3, y=4:6)

期望结果：z = c(2*5+3*6, 3*6, NA)，即c(28, 18, NA)。

原函数的问题分析 #

你编写的sumprod函数无法正常运行，主要有两个问题：

• 循环索引范围错误：1:length(x)-1会被解析为(1:length(x)) - 1，当length(x)=3时得到0,1,2，而非预期的1:2，需改为1:(length(x)-1)。
• 不必要的shift使用：shift(type="lag")的逻辑不符合需求，直接截取后续行计算更直观高效。

修正后的循环实现 #

如果坚持用循环，可修改函数如下：

sumprod <- function(x, y){
  n <- length(x)
  z <- vector("numeric", n) # 预先分配向量长度，提升效率
  for (i in 1:(n-1)){
    # 直接取i+1到末尾的x、y相乘后求和
    z[i] <- sum(x[(i+1):n] * y[(i+1):n], na.rm = FALSE)
  }
  z[n] <- NA # 最后一行无后续行，设为NA
  return(z)
}

# 新增列
dt[, z := sumprod(x, y)]

更高效的data.table向量化实现 #

data.table适合向量化操作，避免循环可大幅提升大数据集的处理效率：

# 一步到位：反向累积求和后移位
dt[, z := shift(cumsum(rev(x*y))[.N:1], 1, type="lead")]

逻辑说明：

1. 计算所有行的x*y乘积；
2. 反转乘积序列后做累积求和，得到从末尾到当前位置的和；
3. 再次反转序列，得到从当前行到末尾的乘积和；
4. 用shift(1, type="lead")将结果向下移位一位，正好对应从i+1行开始的乘积和，最后一行自动为NA。

验证结果 #

执行后查看数据：

dt
#    x y  z
# 1: 1 4 28
# 2: 2 5 18
# 3: 3 6 NA

完全符合预期。

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