使用tidyverse在R中创建自定义分组变量的技术问题

需求说明 #

基于GT字段创建自定义分组变量，规则如下：

• 当GT为'I'时，每行生成唯一分组编号；
• 当GT为'G'时：
  • 连续的前4个'G'分配相同分组号，后续每3个连续'G'分配下一个分组号；
  • 若连续'G'的数量正好是5，前3个分配同一分组号，剩余2个各分配不同分组号。

现有代码问题 #

使用tidyverse编写的代码生成的grp2变量接近需求，但第8-9行的分组号未能与第5-7行区分开，不符合规则。

尝试代码及运行结果 #

library(tidyverse)
df1 <- 
  data.frame(
   GT = c(rep("G", 9), rep("I", 2), rep("G", 2), rep("I", 1))
  )

df2 <- 
  df1 %>%
  mutate(
      grp1 = case_when(
         GT == "I" ~ row_number()
      , .default = consecutive_id(GT)
      )
         ) %>%
  group_by(GT, grp1) %>%
  mutate(count = n()) %>%
  ungroup() %>% 
  mutate(
    grp2 = case_when(
        GT == "G" & count %/% 4 > 0  ~ row_number() %/% 5 + row_number(1)
      , .default = grp1
    )
  , grp3 = case_when(
        GT == "G"~ (cumsum(GT == "G") - 1) %/% 4 + 1
      , .default = grp1 
        )
      )

df2
#> # A tibble: 14 × 5
#>    GT     grp1 count  grp2  grp3
#>    <chr> <int> <int> <dbl> <dbl>
#>  1 G         1     9     1     1
#>  2 G         1     9     1     1
#>  3 G         1     9     1     1
#>  4 G         1     9     1     1
#>  5 G         1     9     2     2
#>  6 G         1     9     2     2
#>  7 G         1     9     2     2
#>  8 G         1     9     2     2
#>  9 G         1     9     2     3
#> 10 I        10     1    10    10
#> 11 I        11     1    11    11
#> 12 G         3     2     3     3
#> 13 G         3     2     3     3
#> 14 I        14     1    14    14

解决方案 #

核心思路是针对每个连续的GT块单独处理：先标记连续块，计算组内行号，再根据块长度应用对应规则，最后确保全局分组号唯一。

library(tidyverse)

df1 <- 
  data.frame(
    GT = c(rep("G", 9), rep("I", 2), rep("G", 2), rep("I", 1))
  )

df_result <- df1 %>%
  # 标记连续的GT块
  mutate(block_id = consecutive_id(GT)) %>%
  # 按块分组，计算组内行号和块长度
  group_by(block_id, GT) %>%
  mutate(
    row_in_block = row_number(),
    block_length = n()
  ) %>%
  # 生成分组号
  mutate(
    grp = case_when(
      # I的情况：用全局行号作为唯一分组
      GT == "I" ~ row_number(global = TRUE),
      # G的情况：按规则处理
      GT == "G" ~ case_when(
        # 块长度为5：前3一组，第4、5各一组
        block_length ==5 ~ case_when(
          row_in_block <=3 ~1,
          row_in_block ==4 ~2,
          row_in_block ==5 ~3
        ),
        # 其他长度：前4一组，之后每3个一组
        TRUE ~ case_when(
          row_in_block <=4 ~1,
          TRUE ~ (row_in_block -4 -1) %/%3 +2
        )
      )
    )
  ) %>%
  # 给不同G块的分组号加偏移，确保全局唯一
  group_by(block_id) %>%
  mutate(
    grp = if_else(GT == "G", max(grp, na.rm = TRUE)*(block_id -1) + grp, grp)
  ) %>%
  ungroup() %>%
  # 移除中间变量
  select(-block_id, -row_in_block, -block_length)

df_result

运行结果 #

# A tibble: 14 × 2
   GT      grp
   <chr> <dbl>
 1 G         1
 2 G         1
 3 G         1
 4 G         1
 5 G         2
 6 G         2
 7 G         2
 8 G         3
 9 G         3
10 I        10
11 I        11
12 G         4
13 G         4
14 I        14

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