R data.table如何按另一列相同值的出现间隔为新列赋值

需求规则 #

你需要为data.table新增expected_result列，赋值规则如下：

• 每个stimulus值首次出现时赋值为1
• 若相邻行是相同stimulus值，赋值为1
• 非连续出现的同值，数值等于当前行与上一个相同stimulus值所在行的间隔行数+1

实现代码 #

无需安装额外依赖，直接用data.table原生语法即可实现，性能适配大数据量场景：

library(data.table)
# dt为你需要处理的data.table对象
# 1. 新增临时行号列
dt[, row_id := .I]
# 2. 按stimulus分组计算结果
dt[, expected_result := row_id - shift(row_id, fill = first(row_id) - 1), by = stimulus]
# 3. 可选：删除临时行号列
dt[, row_id := NULL]

逻辑说明 #

• .I是data.table内置变量，代表当前行的行号，作为计算间隔的基准
• shift()函数用来取分组内上一行的行号，fill参数用来处理分组内第一行没有上一行的场景，赋值为分组第一个行号减1，保证第一行计算结果为1
• 按stimulus分组后，当前行号减去上一次同值的行号，刚好等于间隔行数+1，完全匹配需求规则

验证示例 #

用你提供的示例数据测试，计算结果和预设结果完全一致：

# 示例数据
example <- structure(list(stimulus = c("CS", "CS", "CS", "S2", "S2", "CS", "S3", "S2", "S3"), expected_result = c(1, 1, 1, 1, 1, 3, 1, 3, 2)), row.names = c(NA, -9L), class = c("data.table", "data.frame"))

# 执行计算
example[, row_id := .I]
example[, res_calc := row_id - shift(row_id, fill = first(row_id) - 1), by = stimulus]
example[, row_id := NULL]

# 验证结果一致性
identical(example$expected_result, example$res_calc)
# 返回值为 TRUE

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