为R数据框中id列的非NA连续序列生成唯一组编号

问题说明 #

现有如下R数据框df：

> df
   id value
1   1     1
2   2     3
3   3     2
4  NA     1
5  NA     3
6   8     4
7   9     2
8  10     1
9  NA     1
10 NA     3
11 15     2
12 16     1
13 NA     3
14 NA     4
15 NA     2
16 20     1
17 21     1
18 22     3
19 NA     2
20 NA     1
21 NA     3
22 66     4
23 67     2
24 68     1

需要为id列中每一段连续的非NA值生成唯一组编号（示例变量为id2），中间的NA行对应的id2也设为NA。比如第一段非NA的id值（1、2、3）对应id2=1，后续NA行id2=NA；第二段非NA的id值（8、9、10）对应id2=2，以此类推，最终期望输出如下：

id value id2
1   1     1   1
2   2     3   1
3   3     2   1
4  NA     1  NA
5  NA     3  NA
6   8     4   2
7   9     2   2
8  10     1   2
9  NA     1  NA
10 NA     3  NA
11 15     2   3
12 16     1   3
13 NA     3  NA
14 NA     4  NA
15 NA     2  NA
16 20     1   4
17 21     1   4
18 22     3   4
19 NA     2  NA
20 NA     1  NA
21 NA     3  NA
22 66     4   5
23 67     2   5
24 68     1   5

数据的dput代码：

df <- structure(list(id = c(1, 2, 3, NA, NA, 8, 9, 10, NA, NA, 15, 
16, NA, NA, NA, 20, 21, 22, NA, NA, NA, 66, 67, 68), value = c(1, 
3, 2, 1, 3, 4, 2, 1, 1, 3, 2, 1, 3, 4, 2, 1, 1, 3, 2, 1, 3, 4, 
2, 1)), class = "data.frame", row.names = c(NA, -24L))

解决方案 #

方法1：基础R实现 #

利用cumsum()和is.na()函数标记连续非NA段：

# 生成组编号：当id从NA变为非NA时，组号+1
df$id2 <- cumsum(!is.na(df$id) & c(TRUE, is.na(df$id[-nrow(df)])))
# 将NA行的id2设为NA
df$id2[is.na(df$id)] <- NA

原理：!is.na(df$id) & c(TRUE, is.na(df$id[-nrow(df)]))会在每一段非NA的起始位置生成TRUE，其余位置为FALSE；cumsum()对这些逻辑值求和，就能得到递增的组编号，最后把NA行的组号替换为NA即可。

方法2：tidyverse（dplyr）实现 #

如果习惯使用tidyverse工具链，可以用dplyr的窗口函数：

library(dplyr)

df <- df %>%
  mutate(
    # 标记非NA段的起始行
    start = !is.na(id) & lag(is.na(id), default = TRUE),
    # 累积求和生成组编号
    id2 = cumsum(start),
    # NA行的id2设为NA
    id2 = ifelse(is.na(id), NA, id2)
  ) %>%
  select(-start)  # 移除临时变量

更简洁的写法：

df <- df %>%
  mutate(id2 = cumsum(!is.na(id) & lag(is.na(id), default = TRUE)) %>% 
           replace(is.na(id), NA))

运行上述任意代码后，df的id2列会与期望输出完全一致。

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