R语言中按strata和category两类拆分数据后执行回归的实现方法

基础R实现方案 #

直接通过两次split+嵌套lapply完成需求，无需加载额外包：

# 原始示例数据
n <- 3
strata <- rep(1:4, each=n)
y <- rnorm(n =12)
x <- 1:12
category <- rep(c("A", "B", "C"), times = 4)
df <- cbind.data.frame(y, x, strata, category)

# 第一步：按strata拆分数据为列表
split_strata <- split(df, df$strata)
# 第二步：遍历每个strata子数据框，按category二次拆分后执行回归
reg_res <- lapply(split_strata, function(s_df) {
  split_cat <- split(s_df, s_df$category)
  lapply(split_cat, function(c_df) {
    # 增加样本量判断，避免行数不足回归报错
    if(nrow(c_df) >= 2) {
      lm(y ~ x, data = c_df)
    } else {
      "观测值数量小于2，无法拟合回归"
    }
  })
})

输出结果为嵌套列表，通过reg_res[[strata值]][[category值]]即可提取对应分组的回归结果，例如reg_res[["2"]][["B"]]可获取strata=2、category=B的回归输出。

tidyverse实现方案 #

依赖dplyr和purrr包，输出结构化结果更方便后续批量处理回归统计量：

library(dplyr)
library(purrr)

reg_res_tidy <- df %>%
  # 按两个分组字段直接嵌套生成子数据框，无需两次拆分
  group_by(strata, category) %>%
  nest() %>%
  mutate(
    # 批量执行回归
    model = map(data, ~ if(nrow(.x)>=2) lm(y~x, data = .x) else "观测值不足")
  )

输出为表格格式，每一行对应一个strata+category的唯一分组，可直接筛选对应行提取回归结果，也可通过broom包的tidy/glance函数批量提取系数、拟合优度等指标。

以上两种方案均自动适配不同长度的strata、不同数量的内部category组合，无需手动指定分组维度，可直接用于实际业务数据。

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