统计R数据框行中从首列开始的目标值连续出现次数

需求说明 #

统计R数据框每行中从第一列开始的目标值（此处为1）的连续出现次数，将结果存入新变量count。

示例数据 #

输入数据 #

df <- data.frame(
  A = c(1, 0, 1, 0, 1),
  B = c(0, 0, 1, 1, 1),
  C = c(1, 1, 0, 1, 1),
  D = c(0, 1, 0, 0, 1),
  E = c(1, 0, 1, 0, 0)
)

期望输出 #

df <- data.frame(
  A = c(1, 0, 1, 0, 1),
  B = c(0, 0, 1, 1, 1),
  C = c(1, 1, 0, 1, 1),
  D = c(0, 1, 0, 0, 1),
  E = c(1, 0, 1, 0, 0),
  count = c(1, 0, 2, 0, 4)
)

原代码问题 #

你之前的代码用rle找每行中1的最长连续段，但没限定必须从首列开始，会把行内任意位置的最长连续1算进去，不符合需求。

解决方案 #

这里提供两种精准实现的方式：

方法1：定位第一个非1元素 #

df$count <- apply(df, 1, function(x) {
  first_non_one <- which(x != 1)[1]
  if (is.na(first_non_one)) {
    length(x)  # 整行都是1的情况
  } else if (first_non_one == 1) {
    0  # 第一个元素就不是1的情况
  } else {
    first_non_one - 1  # 从首列到第一个非1元素前的长度
  }
})

方法2：利用累积乘积快速判断 #

df$count <- apply(df, 1, function(x) {
  # cumprod会让第一个非1之后的所有值都变成0，刚好对应从首列开始的连续1范围
  sum(cumprod(x == 1))
})

两种方法都能得到期望的结果，第二种代码更简洁，推荐使用。

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