高效实现按group统计class1与class2所有组合计数的方法

问题背景 #

给定数据集示例：

df <- tibble(group = c(rep(1, 6), rep(2, 6)),
             class1 = c("A", "A", "B", "B", "B", "C", "B", "B", "B", "C", "C", "C"),
             class2 = c("A", "B", "B", "B", "C", "B", "B", "B", "A", "C", "A", "B"))

需求：按group分组，生成class1与class2所有可能组合的计数表格（包含出现次数为0的组合）。

原方案处理千万级数据集时速度极慢（耗时超30分钟）：

output <- df %>% table() %>% as.data.table()

期望输出需包含所有组合，示例如下（注：示例中group字段存在笔误，应为rep(1,9), rep(2,9)）：

output <- tibble(group = c(rep(1, 9), rep(2, 9)),
                 class1 = c(rep("A", 3), rep("B", 3), rep("C", 3),
                            rep("A", 3), rep("B", 3), rep("C", 3)),
                 class2 = rep(c("A", "B", "C"), 6),
                 N = c(1, 1, 0, 0, 2, 1, 0, 1, 0, 0, 0, 0, 1, 2, 0, 1, 1, 1))

高效解决方案 #

方案1：使用data.table（推荐，千万级数据最优） #

data.table的分组操作和交叉连接性能远优于基础函数，适合大数据集：

library(data.table)

# 转换为data.table格式
setDT(df)

# 获取class1和class2的所有唯一取值
class1_vals <- unique(df$class1)
class2_vals <- unique(df$class2)

# 1. 分组统计已有组合的计数；2. 生成每个group下的全组合；3. 左连接补全计数并填充0
output <- df[, .N, by = .(group, class1, class2)] %>%
  .[CJ(group = unique(group), class1 = class1_vals, class2 = class2_vals, sorted = FALSE),
    on = .(group, class1, class2)] %>%
  .[is.na(N), N := 0] %>%
  .[order(group, class1, class2)]

方案2：使用dplyr + tidyr（tidyverse用户友好） #

如果习惯tidyverse语法，可使用count()配合complete()实现，性能优于原方案：

library(dplyr)
library(tidyr)

output <- df %>%
  # 分组统计已有组合的计数
  count(group, class1, class2) %>%
  # 生成所有可能组合，缺失计数填充为0
  complete(group, class1 = unique(class1), class2 = unique(class2), fill = list(N = 0)) %>%
  # 按顺序排序，匹配期望输出格式
  arrange(group, class1, class2)

方案优势说明 #

• 原方案的table()会先构建多维列联表，再转换为数据框，中间过程内存占用大、计算冗余；
• 上述方案直接先分组统计已有组合的计数，再通过交叉连接补全缺失组合，内存占用更低，计算逻辑更高效，处理千万级数据的耗时可压缩至数分钟甚至更短。

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