R语言中遍历DataTable为每个ID构建模型预测第9周销量的方法

嘿，针对你这个按ID单独建模预测第9周Sales的需求，我给你准备了几个可行方案——既有你好奇的for循环写法，也有R圈子里更推荐的分组处理方式，毕竟普通循环在R里确实容易效率拉胯，尤其是数据量大的时候。

先给你写最基础的for循环，满足你想了解R循环写法的需求。这里特意加了对训练数据不足的ID的处理，因为单条数据根本拟合不了glm模型，会报错，所以我们给这类ID设置了默认预测值，你可以按需调整。

# 获取所有唯一ID
unique_ids <- unique(train$ID)

# 初始化test的pred列
test$pred <- NA

# 循环每个ID
for (id in unique_ids) {
  # 提取当前ID的训练数据
  train_subset <- train[train$ID == id, ]
  # 定位当前ID在test中的行
  test_subset_idx <- test$ID == id
  
  # 处理训练数据不足的情况（少于2条无法拟合glm）
  if (nrow(train_subset) < 2) {
    # 可选逻辑：用该ID唯一的Sales值，或者全局Sales均值
    test$pred[test_subset_idx] <- ifelse(nrow(train_subset) == 1, train_subset$Sales, mean(train$Sales, na.rm = TRUE))
    next # 跳过后续建模步骤
  }
  
  # 拟合glm模型
  model <- glm(Sales ~ V1 + V2, data = train_subset)
  # 预测并赋值到test的pred列
  test$pred[test_subset_idx] <- predict(model, newdata = test[test_subset_idx, ], type = "response")
}

如果你习惯用tidyverse这套工具（dplyr、purrr这些），那这个方案会很顺手，代码简洁还易读，而且比普通循环高效不少。核心思路是把每个ID的训练数据嵌套成列表，然后对每个列表元素单独建模，最后关联测试数据做预测。

library(dplyr)
library(purrr)
library(tidyr)

# 第一步：按ID分组嵌套训练数据，并为每个ID拟合模型
train_models <- train %>%
  group_by(ID) %>%
  nest() %>% # 把每个ID的所有行嵌套成一个数据框列表
  mutate(
    model = map(data, function(df) {
      # 处理数据不足的情况
      if (nrow(df) < 2) {
        # 返回一个自定义的"模型"对象，用来存储默认预测值
        list(default_pred = ifelse(nrow(df) == 1, df$Sales, mean(train$Sales, na.rm = TRUE)))
      } else {
        # 正常拟合glm模型
        glm(Sales ~ V1 + V2, data = df)
      }
    })
  )

# 第二步：关联测试数据与模型，完成预测
test_pred <- test %>%
  left_join(train_models, by = "ID") %>% # 把对应ID的模型关联过来
  mutate(
    pred = map2_dbl(model, ID, function(mod, id) {
      # 判断是自定义默认模型还是正规glm模型
      if ("default_pred" %in% names(mod)) {
        mod$default_pred
      } else {
        predict(mod, newdata = cur_data(), type = "response")
      }
    })
  ) %>%
  select(-data, -model) # 移除不需要的嵌套列和模型列

# 最终test_pred就是带预测值的测试表，也可以把pred列合并回原test表

要是你的数据量特别大（比如几十万甚至上百万行），那必须试试data.table，它的分组操作是底层C语言优化过的，速度快到飞起，内存占用也更低。

library(data.table)

# 把数据转换为data.table格式
setDT(train)
setDT(test)

# 第一步：按ID训练模型，存储到列表中
model_list <- train[, {
  if (.N < 2) { # .N是当前分组的行数
    # 处理数据不足的情况，存储默认预测值
    list(model = list(list(default_pred = ifelse(.N == 1, Sales, mean(train$Sales, na.rm = TRUE)))))
  } else {
    # 正常拟合glm模型
    list(model = list(glm(Sales ~ V1 + V2, data = .SD))) # .SD代表当前分组的所有列
  }
}, by = ID]

# 第二步：关联测试数据并完成预测
test[model_list, on = "ID", pred := {
  mod <- i.model[[1]]
  # 判断模型类型，返回对应预测值
  if ("default_pred" %in% names(mod)) {
    mod$default_pred
  } else {
    predict(mod, newdata = .SD, type = "response")
  }
}, by = .EACHI] # .EACHI表示对每个匹配的分组单独处理

补充说明 #

• 对于只有1条训练数据的ID，上面的方案都做了兜底处理，你可以根据业务需求调整默认预测值的逻辑（比如用同行业均值、全局中位数等）。
• 为什么不推荐普通for循环？因为R的普通循环是逐元素解释执行，而分组函数（dplyr的group_by、data.table的by）都是底层优化过的，数据量越大，效率差距越明显。
• 如果你的所有ID都至少有2条训练数据，可以直接去掉那些处理数据不足的判断逻辑，代码会更简洁。

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