如何在R中按nr分组，基于sp=="tr"实现前后计数？

原始数据框 #

df <- data.frame(nr = c(rep("n01", 10),
                        rep("n03", 13),
                        rep("n04", 8),
                        rep("n06", 14),
                        rep("n08", 13),
                        rep("n12", 14),
                        rep("n14", 10)),
                 yr = c(2012:2021, 
                        2010:2022,
                        2013:2020,
                        2010:2023,
                        2011:2023,
                        2009:2022,
                        2011:2020),
                 sp = c(rep(NA, 12),
                        "tr",
                        rep(NA, 27),
                        "tr",
                        rep(NA, 21),
                        "tr",
                        rep(NA, 19)))

需求说明 #

按nr字段分组处理：

• 若分组内存在sp == "tr"的行，以该行作为基准（计数为0），向前行计数为负数，向后行计数为正数；
• 若分组内无符合sp == "tr"的行，count列全部填充NA。

目标数据框（参考） #

df_target <- data.frame(nr = c(rep("n01", 10),
                               rep("n03", 13),
                               rep("n04", 8),
                               rep("n06", 14),
                               rep("n08", 13),
                               rep("n12", 14),
                               rep("n14", 10)),
                        yr = c(2012:2021, 
                               2010:2022,
                               2013:2020,
                               2010:2023,
                               2011:2023,
                               2009:2022,
                               2011:2020),
                        sp = c(rep(NA, 12),
                               "tr",
                               rep(NA, 27),
                               "tr",
                               rep(NA, 21),
                               "tr",
                               rep(NA, 19)),
                        count = c(rep(NA, 10),
                                  -2:10,
                                  rep(NA, 8),
                                  -9:4,
                                  rep(NA, 13),
                                  -4:9,
                                  rep(NA, 10)))

完整解决方案 #

使用dplyr包的分组函数结合位置计算实现：

library(dplyr)

df_result <- df %>%
  group_by(nr) %>%
  mutate(
    count = if (any(sp == "tr", na.rm = TRUE)) {
      # 获取分组内第一个"tr"所在的行位置
      tr_pos <- which(sp == "tr")[1]
      # 用当前行的位置序号减去基准位置，得到计数
      seq_along(sp) - tr_pos
    } else {
      # 无"tr"则填充NA
      rep(NA, n())
    }
  ) %>%
  ungroup()

代码说明 #

1. group_by(nr)：按nr字段分组；
2. any(sp == "tr", na.rm = TRUE)：判断分组内是否存在有效"tr"值（排除NA干扰）；
3. which(sp == "tr")[1]：定位分组内第一个"tr"的行位置（数据中每组仅存在一个"tr"，无需处理多个的情况）；
4. seq_along(sp) - tr_pos：生成从基准行向前、向后的计数序列，基准行计数为0，向前依次为-1、-2...，向后依次为1、2...；
5. ungroup()：取消分组，回归普通数据框格式。

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