按用户组为连续n个非NA值生成枚举序号列的实现方法

需求说明 #

按用户分组，识别出连续长度≥n的非NA值序列，为这些序列中的值添加连续递增的枚举序号；对于长度不足n的连续非NA序列，对应的位置标记为NA。

示例数据 #

library(dplyr)

n <- 4

# 示例数据框
df <- tibble(
  user_id = c(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2),
  date = as.Date('2024-01-01') + 0:19,
  value1 = c(NA, 5, 6, 7, 8, NA, NA, 2, 3, 4, 5, NA, 1, 2, 3, NA, 1, 2, 3, 4)
)

解决方案 #

使用dplyr包通过分组、标记连续块、筛选有效序列三个步骤实现需求：

df_result <- df %>%
  # 按用户分组
  group_by(user_id) %>%
  # 标记连续非NA值的块：每遇到一个NA，块ID递增
  mutate(block_id = cumsum(is.na(value1))) %>%
  # 按用户+块ID分组，计算每个块的非NA值数量
  group_by(user_id, block_id) %>%
  mutate(block_size = sum(!is.na(value1))) %>%
  # 回到用户分组，生成目标序号
  group_by(user_id) %>%
  mutate(
    # 标记当前行是否属于有效序列（非NA且所在块长度≥n）
    is_valid = !is.na(value1) & block_size >= n,
    # 对有效行累加计数，无效行设为NA
    expected_result = ifelse(is_valid, cumsum(is_valid), NA)
  ) %>%
  # 清理临时列
  select(-block_id, -block_size, -is_valid) %>%
  ungroup()

# 查看结果
print(df_result)

结果验证 #

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

# A tibble: 20 × 4
   user_id date       value1 expected_result
     <dbl> <date>      <dbl>           <dbl>
 1       1 2024-01-01     NA              NA
 2       1 2024-01-02      5               1
 3       1 2024-01-03      6               2
 4       1 2024-01-04      7               3
 5       1 2024-01-05      8               4
 6       1 2024-01-06     NA              NA
 7       1 2024-01-07     NA              NA
 8       1 2024-01-08      2               5
 9       1 2024-01-09      3               6
10       1 2024-01-10      4               7
11       1 2024-01-11      5               8
12       1 2024-01-12     NA              NA
13       2 2024-01-13      1              NA
14       2 2024-01-14      2              NA
15       2 2024-01-15      3              NA
16       2 2024-01-16     NA              NA
17       2 2024-01-17      1               1
18       2 2024-01-18      2               2
19       2 2024-01-19      3               3
20       2 2024-01-20      4               4

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