高效实现大数据集中特定条件成对比较求和与新列创建

问题背景 #

我有如下数据集：

dat1 <- read.table(text = "
  nodepair 3 4 5
1    A6_A1 2 5 1
2    A6_A2 2 5 1
3    A6_A3 2 5 1
4    AL_A1 1 0 0
5     D_A6 0 3 0
6     F_A1 1 0 1
7      H_D 0 0 2
8      H_H 0 0 2 
", header = TRUE, check.names = FALSE)

需要编写高效代码，创建新列对成对变量进行条件求和：

• b类型列：当成对变量中第二个变量为0时，对两个变量求和，否则为NA
• c类型列：当成对变量中第一个变量为0时，对两个变量求和，否则为NA

期望输出如下：

dat2 <- read.table(text = "
  nodepair 3 4 5 3-4b 3-4c 3-5b 3-5c 4-5b 4-5c
1    A6_A1 2 5 1   NA   NA   NA   NA   NA   NA
2    A6_A2 2 5 1   NA   NA   NA   NA   NA   NA
3    A6_A3 2 5 1   NA   NA   NA   NA   NA   NA
4    AL_A1 1 0 0    1   NA    1   NA    0    0
5     D_A6 0 3 0   NA    3    0    0    3   NA
6     F_A1 1 0 1    1   NA   NA   NA   NA    1
7      H_D 0 0 2    0    0   NA    2   NA    2
8      H_H 0 0 2    0    0   NA    2   NA    2  
", header = TRUE, check.names = FALSE)

我目前用以下代码实现，但大数据集下效率极低，寻求更高效的实现方式：

dat1 <- dat1 %>%
  mutate('3-4b' = case_when(`4` == 0 ~ as.integer(rowSums(across(c(`3`,`4`)))))) %>%
  mutate('3-4c' = case_when(`3` == 0 ~ as.integer(rowSums(across(c(`3`,`4`)))))) %>%
  mutate('3-5b' = case_when(`5` == 0 ~ as.integer(rowSums(across(c(`3`,`5`)))))) %>%
  mutate('3-5c' = case_when(`3` == 0 ~ as.integer(rowSums(across(c(`3`,`5`)))))) %>%
  
  mutate('4-5b' = case_when(`5` == 0 ~ as.integer(rowSums(across(c(`4`,`5`)))))) %>%
  mutate('4-5c' = case_when(`4` == 0 ~ as.integer(rowSums(across(c(`4`,`5`)))))) 

高效解决方案 #

方案1：dplyr批量处理（减少重复计算） #

通过生成所有数值列的两两组合，批量计算新列，避免多次调用mutate：

library(dplyr)

# 提取数值列名称
num_cols <- names(dat1)[-1]
# 生成所有两两组合
pair_combs <- t(combn(num_cols, 2))

# 批量生成新列
dat1_optimized <- dat1 %>%
  bind_cols(
    pmap_dfc(as.data.frame(pair_combs), function(x, y) {
      # 计算b类型列：第二个变量为0时求和
      col_b <- case_when(!!sym(y) == 0 ~ as.integer(!!sym(x) + !!sym(y)), TRUE ~ NA_integer_)
      # 计算c类型列：第一个变量为0时求和
      col_c <- case_when(!!sym(x) == 0 ~ as.integer(!!sym(x) + !!sym(y)), TRUE ~ NA_integer_)
      # 返回命名后的列
      tibble(!!paste0(x, "-", y, "b") := col_b, !!paste0(x, "-", y, "c") := col_c)
    })
  )

方案2：data.table（极致性能，适合超大数据集） #

data.table的操作在百万级以上数据量下性能远超dplyr，适合处理大规模数据集：

library(data.table)

setDT(dat1)
num_cols <- names(dat1)[-1]
pair_combs <- t(combn(num_cols, 2))

# 循环生成新列
for (i in seq(nrow(pair_combs))) {
  x <- pair_combs[i, 1]
  y <- pair_combs[i, 2]
  # 计算b列：用fifelse替代case_when提升性能
  dat1[, paste0(x, "-", y, "b") := fifelse(get(y) == 0, as.integer(get(x) + get(y)), NA_integer_)]
  # 计算c列
  dat1[, paste0(x, "-", y, "c") := fifelse(get(x) == 0, as.integer(get(x) + get(y)), NA_integer_)]
}

方案说明 #

• 两种方案都通过批量处理成对组合避免了重复的列操作，减少了不必要的计算开销
• 直接用x + y替代rowSums(across(c(x,y)))，逻辑更简洁，计算速度更快
• data.table的fifelse比dplyr的case_when性能更优，数据量越大优势越明显

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