如何按label分组计算val列累积和并添加末次累积值列

问题背景 #

现有如下R矩阵：

x <- cbind(label = c(1,1,1,-1,-1,0,0,-1,1,1),
            val = c(1,3,7,5,4,2,6,8,5,1))
# 输出结果：
#       label val
#  [1,]     1   1
#  [2,]     1   3
#  [3,]     1   7
#  [4,]    -1   5
#  [5,]    -1   4
#  [6,]     0   2
#  [7,]     0   6
#  [8,]    -1   8
#  [9,]     1   5
# [10,]     1   1

需要实现两个目标：

1. 根据label的连续序列计算val的累积和：
   • 连续1时，计算正向累积和；
   • 连续-1时，计算负向累积和（即每个val乘以-1后累积）；
   • label为0时，对应累积和设为0；
2. 新增last_cum_val列，仅在每组连续非0label序列的最后一行记录该组的最终累积和，其余位置填NA。

期望输出结果：

res_x <- cbind(x,
                cum_val = c(1,4,11,-5,-9,0,0,-8,5,6),
                last_cum_val = c(NA,NA,11,NA,-9,NA,NA,-8,NA,6))
# 输出结果：
#       label val cum_val last_cum_val
#  [1,]     1   1       1           NA
#  [2,]     1   3       4           NA
#  [3,]     1   7      11           11
#  [4,]    -1   5      -5           NA
#  [5,]    -1   4      -9           -9
#  [6,]     0   2       0           NA
#  [7,]     0   6       0           NA
#  [8,]    -1   8      -8           -8
#  [9,]     1   5       5           NA
# [10,]     1   1       6            6

注：原示例中第8行last_cum_val写为8，应为笔误，按需求逻辑正确值为-8。

解决方案 #

方法1：使用data.table包 #

library(data.table)

# 转换为data.table格式
dt <- as.data.table(x)

# 生成分组标识：将连续相同的非0label归为一组，0单独成组
dt[, group := rleid(label != 0, label)]

# 计算cum_val列
dt[, cum_val := ifelse(label == 0, 0, cumsum(val * label)), by = group]

# 计算last_cum_val列：仅每组最后一行（非0label组）记录最终累积和
dt[, last_cum_val := ifelse(.I == max(.I) & label != 0, cum_val, NA), by = group]

# 转换回矩阵格式
res_x <- as.matrix(dt)
print(res_x)

方法2：使用dplyr包 #

library(dplyr)

# 转换为data.frame格式
df <- as.data.frame(x)

# 生成分组标识
df <- df %>%
  mutate(group = rleid(label != 0, label))

# 分组计算cum_val和last_cum_val
df <- df %>%
  group_by(group) %>%
  mutate(
    cum_val = ifelse(label == 0, 0, cumsum(val * label)),
    last_cum_val = ifelse(row_number() == n() & label != 0, cum_val, NA)
  ) %>%
  ungroup()

# 转换回矩阵格式
res_x <- as.matrix(df)
print(res_x)

内容的提问来源于stack exchange，提问作者mr.T