R中sum() over(partition by)等价实现：含NA分组求和问题解决

原始数据 #

df <- data.frame(
  id=c("1", "1", "2", "2", "3", "4"),
  tube_placement=c("2020-01-01", "2020-01-10", "2020-01-01", "2020-01-15", "2020-01-01", "" ),
  tube_removal = c("2020-01-02", "2020-01-12", "2020-01-02", "", "2020-01-02", ""),
  attempts = c(1, 2, 1, "", 1, "")
)
df[df==""] <- NA

df$attempts <- as.numeric(df$attempts)

预期结果 #

id  tube_placement  tube_removal    attempts    total_attempts
1   2020-01-01      2020-01-02      1           3
1   2020-01-10      2020-01-12      2           3
2   2020-01-01      2020-01-02      1           1
2   2020-01-15      NA              NA          1
3   2020-01-01      2020-01-02      1           1
4   NA              NA              NA          NA  

尝试的代码 #

df1 <- df %>% group_by(id) %>%
  mutate(total_attempts = sum(attempts))

问题原因 #

原代码中sum(attempts)默认会因为组内存在NA值而返回NA（比如id=2的组）；同时对于id=4这种组内所有attempts都是NA的情况，直接用sum(attempts, na.rm=TRUE)会得到0，不符合预期的NA结果。

改进后的代码 #

通过all(is.na(attempts))判断组内是否全为NA，再结合sum(attempts, na.rm=TRUE)实现需求：

library(dplyr)

df1 <- df %>% 
  group_by(id) %>%
  mutate(
    total_attempts = if (all(is.na(attempts))) {
      NA_real_
    } else {
      sum(attempts, na.rm = TRUE)
    }
  ) %>%
  ungroup()

也可以用case_when实现相同逻辑：

df1 <- df %>% 
  group_by(id) %>%
  mutate(
    total_attempts = case_when(
      all(is.na(attempts)) ~ NA_real_,
      TRUE ~ sum(attempts, na.rm = TRUE)
    )
  ) %>%
  ungroup()

运行后即可得到符合预期的total_attempts列。

内容的提问来源于stack exchange，提问作者Prajwal Mani Pradhan