如何用data.table替代循环，基于行间差值为数据分配分组变量？

需求是根据Value字段当前行与上一行的差值分配组标识：当差值≤设定的阈值critdiff时，沿用前一行的组号；当差值>阈值时，组号加1。你已经用for循环实现了需求，这里提供data.table的无循环高效实现方案。

原for循环实现代码 #

require(data.table)

dat <- data.table(ID=1:10, Value=c(1,20,36,234,321,553,556,579,588,701))
dat[1,Group:=1]
critdiff <- 100
n <- dat[,.N]
for(i in 2:n){
    if(dat[i,Value]-dat[i-1,Value]<=critdiff){
    dat[i,Group:=dat[i-1,Group]]
    }else{
    dat[i,Group:=dat[i-1,Group+1]]
    }
}
dat[]

data.table无循环实现方案 #

利用shift()函数获取上一行的Value值，结合cumsum()累加阈值触发的次数，快速生成组标识：

require(data.table)

dat <- data.table(ID=1:10, Value=c(1,20,36,234,321,553,556,579,588,701))
critdiff <- 100

# 计算组号
dat[, Group := 1 + cumsum(Value - shift(Value, type = "lag", fill = Value[1]) > critdiff)]

# 查看结果
dat[]

代码解释 #

• shift(Value, type = "lag", fill = Value[1])：生成Value的滞后向量（上一行的值），第一行无前置行，用自身值填充，确保第一行差值为0。
• Value - shift(...) > critdiff：生成逻辑向量，当当前行与上一行的差值超过阈值时为TRUE，否则为FALSE。
• cumsum()对逻辑向量累加（TRUE视为1，FALSE视为0），每次触发阈值就累加1，最后加初始组号1，得到最终的组标识。

这个方案完全采用向量化操作，避免了循环的性能损耗，在处理大规模数据时优势尤为明显，结果与原for循环完全一致。

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