R语言data.table中基于ID的条件累计扣减与余额生成实现

核心逻辑 #

Balance的计算不需要逐行循环，利用data.table向量化计算特性即可实现：

• 以首行Total值为初始扣减基准
• 仅对ID == "A"的行计入扣减量，ID != "A"的行扣减量按0计算，自然继承前序累计值、不做扣减
• 对有效扣减量做累计求和，用初始基准值减去累计扣减量即可得到每行的Balance值

完整可运行代码 #

library(data.table)
dt <- fread('
   Year   ID   Total      Quantity
   2017   A     2000        100
   2016   A     2000        200 
   2015   A     2000        50
   2014   B     2000        100
   2013   B     2000        300
')

# 生成Balance列
dt[, Balance := Total[1] - cumsum(ifelse(ID == "A", Quantity, 0))]

针对超大数据集可替换为性能更高的写法，用data.table自带的快速条件判断函数fifelse替代基础函数ifelse：

dt[, Balance := Total[1] - cumsum(fifelse(ID == "A", Quantity, 0))]

结果验证 #

运行代码后dt的输出与目标数据集want完全一致：

Year ID Total Quantity Balance
1: 2017  A  2000      100    1900
2: 2016  A  2000      200    1700
3: 2015  A  2000       50    1650
4: 2014  B  2000      100    1650
5: 2013  B  2000      300    1650

注意事项 #

• 累计计算严格按照data.table当前的行顺序从上到下执行，若业务逻辑需要按特定顺序（如按年份升序/降序）累计，需先用setorder()调整行序再计算
• 若不同行的Total值存在差异、需以每行自身的Total为扣减基准，可将公式调整为dt[, Balance := Total - cumsum(fifelse(ID == "A", Quantity, 0))]

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