基于起止索引向量为DataFrame分配分组ID

现有如下DataFrame：

df <- data.frame(X = LETTERS[1:20], 
                 Y = paste0("abc_", 1:20))
# 输出示例：
#  X     Y
#1 A abc_1
#2 B abc_2
#3 C abc_3
#4 D abc_4
#5 E abc_5
#6 F abc_6
# ...

需要根据两个整数向量为每行分配分组ID：

start_ix <- c(2, 5, 8, 10, 15, 18)
end_ix   <- c(4, 7, 9, 13, 17, 19)

规则：

• 第1组对应行索引2-4，第2组对应5-7，以此类推
• 不在任何起止区间内的行，分组ID设为NA

期望输出如下：

df_want <- structure(list(X = c("A", "B", "C", "D", "E", "F", "G", "H", 
"I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T"), 
    Y = c("abc_1", "abc_2", "abc_3", "abc_4", "abc_5", "abc_6", 
    "abc_7", "abc_8", "abc_9", "abc_10", "abc_11", "abc_12", 
    "abc_13", "abc_14", "abc_15", "abc_16", "abc_17", "abc_18", 
    "abc_19", "abc_20"), grp = c(NA, 1, 1, 1, 2, 2, 2, 3, 3, 
    4, 4, 4, 4, NA, 5, 5, 5, 6, 6, NA)), row.names = c(NA, -20L
), class = "data.frame")

# 输出示例：
#    X      Y grp
# 1  A  abc_1  NA
# 2  B  abc_2   1
# 3  C  abc_3   1
# 4  D  abc_4   1
# 5  E  abc_5   2
# 6  F  abc_6   2
# 7  G  abc_7   2
# 8  H  abc_8   3
# 9  I  abc_9   3
# 10 J abc_10   4
# 11 K abc_11   4
# 12 L abc_12   4
# 13 M abc_13   4
# 14 N abc_14  NA
# 15 O abc_15   5
# 16 P abc_16   5
# 17 Q abc_17   5
# 18 R abc_18   6
# 19 S abc_19   6
# 20 T abc_20  NA

要求优先使用base R实现，也欢迎其他R包的解决方案作为参考。

Base R 实现 #

方法1：利用findInterval快速映射区间 #

# 生成每行的索引
row_ix <- seq_len(nrow(df))
# 通过区间匹配分组ID
grp <- findInterval(row_ix, c(1, end_ix))
# 过滤掉不在指定区间的行，设为NA
grp[!(row_ix %in% unlist(mapply(seq, start_ix, end_ix)))] <- NA

# 合并到原DataFrame
df$grp <- grp

方法2：循环赋值（直观易懂） #

# 初始化grp列为NA
df$grp <- NA
# 遍历每个分组的起止索引，批量赋值分组ID
for(i in seq_along(start_ix)){
  df$grp[start_ix[i]:end_ix[i]] <- i
}

其他R包实现 #

dplyr 方案 #

library(dplyr)

# 构建分组区间对照表
grp_df <- tibble(
  start = start_ix,
  end = end_ix,
  grp = seq_along(start_ix)
)

# 通过行索引匹配区间并合并
df <- df %>%
  mutate(row_ix = row_number()) %>%
  left_join(grp_df, join_by(between(row_ix, start, end))) %>%
  select(-row_ix)

data.table 方案 #

library(data.table)

# 转换为data.table格式并添加行索引
setDT(df)[, row_ix := .I]
# 构建分组区间对照表
grp_dt <- data.table(start = start_ix, end = end_ix, grp = seq_along(start_ix))

# 区间匹配并赋值分组ID
df <- df[grp_dt, on = .(row_ix >= start, row_ix <= end), grp := i.grp]
# 移除临时行索引列
df[, row_ix := NULL]

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