如何按分组及自定义条件实现R语言中的计数器重置

1. 输入数据集 #

id = c("A","A","A","A","A","B", "B", "B", "B")
result = c(1,1,0,1,1,0,1,0,1)
my_data = data.frame(id, result)

2. 计数器规则 #

• 按唯一id分组独立计算计数器
• 每组首个值若为1，计数器初始为1；否则为0
• 遇到result=1时，计数器加1
• 遇到result=0时，计数器重置为0
• 首次遇到1前，计数器保持0；每次遇到1后重新开始递增计数

3. 预期输出 #

id result counter
1  A      1       1
2  A      1       2
3  A      0       0
4  A      1       1
5  A      1       2
6  B      0       0
7  B      1       1
8  B      0       0
9  B      1       1

4. 尝试过的代码 #

# 未按id分组，全局生成计数器
my_data$counter = unlist(lapply(split(my_data$results, c(0, cumsum(abs(diff(!my_data$results == 1))))), function(x) (x[1] == 1) * seq(length(x))))

# 按id分组，但逻辑较难理解
my_data$counter = ave(my_data$results, my_data$id, FUN = function(x){ tmp<-cumsum(x);tmp-cummax((!x)*tmp)})

5. 问题与需求 #

上述两段代码逻辑抽象难以理解，需要实现一个通用函数，支持自定义规则，例如：

• result="AAA"：计数器重置为0
• result="BBB"：计数器+1
• result="CCC"：计数器+2
• result="DDD"：计数器-1

现有代码逻辑解析 #

第一段代码（全局分组） #

这段代码未考虑id分组，核心是把连续的result状态（是否为1）拆分成组：

1. !my_data$results == 1：将result=1转为FALSE，其他值转为TRUE
2. diff(...)：计算相邻元素的状态差异，状态变化时返回1或-1，取绝对值后统一为1
3. cumsum(abs(diff(...)))：生成分组标识，每遇到状态变化就加1，前面补0保证分组长度与原数据一致
4. split(...)：按分组拆分数据，对每个分组：若首个元素是1则生成递增序列，否则全0，最后合并结果。但因为没按id分组，不符合需求。

第二段代码（按id分组） #

这段代码利用向量化操作实现分组计数，逻辑绕但高效：

1. tmp <- cumsum(x)：对每个id组内的result累加求和
2. (!x)*tmp：当result=0时保留当前累加值，否则为0
3. cummax((!x)*tmp)：取到当前位置为止，所有result=0时的最大累加值
4. tmp - cummax(...)：用当前累加值减去之前最大的0点累加值，实现重置后重新计数，最终得到符合要求的结果。

通用自定义计数器函数实现 #

方案1：用dplyr + purrr（简洁高效） #

该方案利用tidyverse工具链，逻辑清晰且易于扩展：

library(dplyr)
library(purrr)

# 定义自定义计数器函数
custom_counter <- function(data, id_col, value_col, rules) {
  # 默认操作：无匹配规则时保持计数器不变
  default_op <- function(current_counter) current_counter
  
  data %>%
    group_by({{id_col}}) %>%
    mutate(
      counter = accumulate({{value_col}}, .init = 0, function(prev, curr) {
        # 获取当前值对应的操作，无匹配则用默认
        op <- rules[[as.character(curr)]] %||% default_op
        # 执行操作并返回新计数器值
        if (is.function(op)) {
          op(prev)
        } else {
          if (op == 0) {
            0  # 重置为0
          } else {
            prev + op  # 加减操作
          }
        }
      }) %>% tail(-1)  # 移除初始的0值
    ) %>%
    ungroup()
}

应用原需求规则

# 定义原需求规则：1则+1，0则重置为0
original_rules <- list(
  "1" = function(x) x + 1,
  "0" = function(x) 0
)

# 生成结果
custom_counter(my_data, id, result, original_rules)

应用自定义规则示例

# 构造测试数据
test_data <- data.frame(
  id = c("X", "X", "X", "Y", "Y", "Y"),
  result = c("AAA", "BBB", "CCC", "BBB", "DDD", "AAA")
)

# 定义自定义规则
custom_rules <- list(
  "AAA" = function(x) 0,
  "BBB" = function(x) x + 1,
  "CCC" = function(x) x + 2,
  "DDD" = function(x) x - 1
)

# 生成结果
custom_counter(test_data, id, result, custom_rules)

输出结果：

# A tibble: 6 × 3
  id    result counter
  <chr> <chr>    <dbl>
1 X     AAA          0
2 X     BBB          1
3 X     CCC          3
4 Y     BBB          1
5 Y     DDD          0
6 Y     AAA          0

方案2：Base R实现（无依赖） #

如果不想依赖tidyverse包，可用Base R实现，逻辑更直观：

# 定义%||%运算符：若x为NULL则返回y，否则返回x
`%||%` <- function(x, y) if (is.null(x)) y else x

custom_counter_base <- function(data, id_col, value_col, rules) {
  id_vec <- data[[id_col]]
  value_vec <- data[[value_col]]
  counter <- numeric(nrow(data))
  
  # 默认操作：保持计数器不变
  default_op <- function(x) x
  
  # 遍历每个唯一id分组处理
  for (id in unique(id_vec)) {
    idx <- which(id_vec == id)
    current_counter <- 0
    for (i in seq_along(idx)) {
      curr_val <- as.character(value_vec[idx[i]])
      op <- rules[[curr_val]] %||% default_op
      
      if (is.function(op)) {
        current_counter <- op(current_counter)
      } else {
        current_counter <- if (op == 0) 0 else current_counter + op
      }
      
      counter[idx[i]] <- current_counter
    }
  }
  
  data$counter <- counter
  return(data)
}

# 用法同dplyr版本，例如：
# custom_counter_base(my_data, "id", "result", original_rules)

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