【进阶】Next N rows when condition is TRUE

问

题引入

本期的问题来自于stackoverflow.com，由于大猫实在想不出简洁的翻译，想来想去还是原标题最能描述问题，所以干脆直接借用。如果硬要翻译的话，大概就是“当某条件成立时，找到这个观测后N行的观测”。

举个例子吧！

在这个数据集中，我们希望每当condition=1时，就标记出它接下来的2行（这里N=2了）。例如，在以上数据集第4行的condition是1， 那么我们能够标记出第5行以及第6行。又由于第6行的分组从a变成了b，所以只有第5行被标记了出来。具体小伙伴们一看desireOutcome的输出就明白了。

看起来似乎无从下手？其实，要实现这一步只需要一行代码哦。

附：生成样例数据集的文件：

# dt是样例数据集，一共有15行。

# a是分组变量；

dt <- data.table(a = rep(c("a", "b", "c"), each = 5))

# condition是条件；desireOutcome是希望获得的结果

dt[, condition := as.numeric(.I %% 4 == 0)]

本文需要用到data.table包！

步

骤分解

我们先把这一行优雅的代码放上来：

dt[, desiredOutcome := Reduce('+', shift(condition, 0:2, fill = 0)), by = a]

现在我们逐一分析这一行代码。

shift函数

它能够对向量进行lag与lead操作。参数0:2的意思是分别滞后0期、1期、2期。参数fill的意思是对于leading missing value，使用0进行补齐。需要注意shift最后输出的是一个list，因为我们生成了三个拥有不同滞后期的向量。例如，如果我们有个向量

v <- c(1, 2, 3, 4, 5)

那么使用下面代码

shift(v, 0:2, fill =0)

生成的结果就是这样：

对于原数据集，如果我们删掉Reduce函数，只保留shift函数

dt[, shift(condition, 0:2, fill = 0), by = a]

那么生成的结果就是这样：

其中，V1-V3分别表示对condition变量滞后0，1，2期的结果

Reduce函数

重点来了！在使用shift函数后，我们实际上生成了三个向量，第一个向量只有条件成立时才为1， 第二个向量条件成立后的“滞后一期”才为1， 第三个向量只有条件成立后的“滞后两期”才为1。于是很自然的，如果我们能将这三个向量相加，那么所有符合要求的行就都是1， 不符合的就都是0了。关键问题在于如果给我们一个list，使用什么方法能够把list的每个元素“一一对应”地加总呢？这时我们就需要用到Reduce函数。它的参数“+”相当于把上图的V1-V3列进行加总并生成新的变量。

综上，完整代码就是:

dt[, desiredOutcome := Reduce('+', shift(condition, 0:2, fill = 0)), by = a]

拓

展

如果想了解更多，可以看stackoverflow.com上的原文：

http://stackoverflow.com/questions/36766452/r-data-table-find-next-n-rows-when-condition-is-true