在R中高效统计指定列排除NA后每行1的数量

需求说明 #

为数据框创建新列X11，根据指定列（示例为X1、X2、X3、X4，后续需扩展至7个变量）的NA数量，统计非NA列中1的个数：

• 若有1个NA：统计剩余3个非NA列中1的数量
• 若有2个NA：统计剩余2个非NA列中1的数量
• 若有3个NA：判断剩余1个非NA列是否为1（是则计1，否则0）
• 若全为NA：结果为0

数据生成代码 #

df <- data.frame(replicate(10,sample(0:2, 10, rep=TRUE)))
df <- replace(df, df == 0, NA)

数据样例 #

X1 X2 X3 X4 X5 X6 X7 X8 X9 X10
1   1  1 NA  1 NA NA NA  1  1   2
2  NA  1  1 NA  2 NA  2  2 NA   1
3   1 NA  1  1 NA NA  1  2 NA   1
4   2  2  2  1  1  2  1 NA  2   2
5  NA  2 NA  2 NA  2  1 NA  1   1
6   2  2  1  1  2 NA  1  2  1   1
7   1  2 NA NA  2  1  1 NA NA   1
8   2  2 NA NA  1 NA NA  2 NA   1
9   1 NA  1  2  2  1  2 NA NA   1
10 NA  2  1 NA NA NA NA  2  2  NA

期望输出 #

X1 X2 X3 X4 X5 X6 X7 X8 X9 X10 X11
1   1  1 NA  1 NA NA NA  1  1   2   3
2  NA  1  1 NA  2 NA  2  2 NA   1   2
3   1 NA  1  1 NA NA  1  2 NA   1   3
4   2  2  2  1  1  2  1 NA  2   2   1
5  NA  2 NA  2 NA  2  1 NA  1   1   0
6   2  2  1  1  2 NA  1  2  1   1   2
7   1  2 NA NA  2  1  1 NA NA   1   1
8   2  2 NA NA  1 NA NA  2 NA   1   0
9   1 NA  1  2  2  1  2 NA NA   1   2
10 NA  2  1 NA NA NA NA  2  2  NA   1

当前问题 #

现有代码通过枚举所有变量组合实现，效率极低，无法处理30万条观测数据，也难以扩展到7个变量。

高效解决方案 #

方案1：tidyverse 向量化实现 #

利用rowSums和across的向量化操作，无需枚举组合，直接统计非NA列中1的个数，再处理全NA的情况：

library(dplyr)

# 指定目标变量
vars <- c("X1", "X2", "X3", "X4")

df <- df %>%
  mutate(
    # 统计非NA列中1的数量
    X11 = rowSums(across(vars, ~.x == 1), na.rm = TRUE),
    # 当所有指定列都是NA时，将X11设为0
    X11 = ifelse(rowSums(is.na(across(vars))) == length(vars), 0, X11)
  )

方案2：Base R 高效实现 #

使用apply函数结合向量判断，同样避免枚举：

vars <- c("X1", "X2", "X3", "X4")

df$X11 <- apply(df[vars], 1, function(row) {
  if (all(is.na(row))) {
    0
  } else {
    sum(row == 1, na.rm = TRUE)
  }
})

说明 #

两种方案均为向量化操作，处理30万条数据速度极快，且可直接扩展到7个变量（只需修改vars向量即可），完全符合需求。

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