如何将R语言数据框分组扩展至最大分组的行数？

原始数据 #

df <- structure(list(ID = c(1L, 1L, 2L, 3L, 3L, 3L), col1 = c("A", 
"B", "O", "U", "L", "R")), class = "data.frame", row.names = c(NA, 
-6L))

需求说明 #

将数据框中每个ID分组扩展至数据集中最大分组的行数（本例中最大行数为3），分组内行数不足的部分用NA填充。

解决方案（tidyverse 版本） #

library(dplyr)
library(tidyr)

# 计算所有分组的最大行数
max_row_count <- df %>% count(ID) %>% pull(n) %>% max()

# 分组扩展并填充NA
expanded_df <- df %>%
  group_by(ID) %>%
  mutate(row_idx = row_number()) %>%  # 给分组内每行编序号
  complete(row_idx = 1:max_row_count) %>%  # 扩展到最大行数，自动补NA
  ungroup() %>%
  select(-row_idx)  # 移除辅助列

# 查看结果
print(expanded_df)

输出结果：

# A tibble: 9 × 2
     ID col1 
  <int> <chr>
1     1 A    
2     1 B    
3    NA NA   
4     2 O    
5    NA NA   
6    NA NA   
7     3 U    
8     3 L    
9     3 R

解决方案（data.table 版本，适合大数据） #

library(data.table)

setDT(df)
max_row_count <- df[, .N, by = ID][, max(N)]

# 按分组扩展行数，自动补NA
expanded_dt <- df[, .SD[1:max_row_count], by = ID]
# 对col1为NA的行，将ID也设为NA
expanded_dt[is.na(col1), ID := NA]

# 查看结果
print(expanded_dt)

逻辑说明 #

1. 先统计每个分组的行数，找到全局最大行数，确定每个分组需要扩展的目标行数。
2. 给每个分组内的行添加序号，利用complete()（tidyverse）或直接索引.SD（data.table）生成目标行数的结构，缺失的行自动填充NA。
3. 最后调整格式，得到符合需求的结果。

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