如何在R data.table中使用动态n值实现分组lag/lead偏移计算

错误原因 #

data.table的shift()函数的n参数仅支持传入单个整数，传入shift_by列向量时，函数会识别为需要生成对应长度的多列偏移结果，不会逐行匹配偏移量，因此返回结果全为NA。

高效解决方案 #

利用分组内的行索引逻辑做向量化取值，无需逐行循环，性能完全满足大数据量场景要求，代码如下：

library(data.table)
df <- data.table(user = c('a', 'a', 'a', 'b', 'b')
                 , spend = 1:5
                 , shift_by = c(1,1,2,1,1)
                 )

# 按user分组实现动态偏移
df[, spend_shifted := {
  # 计算目标取值的全局行号
  target_idx <- .I - shift_by
  # 偏移后超出分组起始行的位置设为NA
  target_idx[target_idx < .I[1]] <- NA_integer_
  # 按索引取值
  spend[target_idx]
}, by = user]

运行后得到的结果和预期完全一致：

user spend shift_by spend_shifted
1:    a     1        1            NA
2:    a     2        1             1
3:    a     3        2             1
4:    b     4        1            NA
5:    b     5        1             4

补充说明 #

如果业务场景允许shift_by取0（取当前行值）或者负值（向前偏移，即lead逻辑），只需要调整target_idx的合法性判断规则即可，整体运算逻辑不需要修改。

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