如何高效重排R语言数据框行以生成重复序列？

需求说明 #

现有数据集，其中每个name值会重复若干次，需要将原本连续重复的name重排为循环序列（如a-b-c循环）。

示例输入 #

test_data <- data.frame(
  name = rep(c("a", "b", "c"), each = 3),
  value = 1:9
)

输出：

name value
1    a     1
2    a     2
3    a     3
4    b     4
5    b     5
6    b     6
7    c     7
8    c     8
9    c     9

期望输出 #

name value
1    a     1
2    b     4
3    c     7
4    a     2
5    b     5
6    c     8
7    a     3
8    b     6
9    c     9

现有实现及性能瓶颈 #

已实现的方案通过拆分数据集、生成排序ID再合并排序，可行但在百万级数据上性能一般：

my_test <- function() {
  split_data <- split(test_data, ~ name)
  
  list_data <- lapply(seq_along(split_data), function(x) {
    new_id <- x + (1:nrow(split_data[[x]]) - 1) * length(split_data)
    split_data[[x]]$id2 <- new_id
    return(split_data[[x]])
  })
  
  out <- do.call(rbind, list_data)
  out <- out[order(out$id2), ]
  out$id2 <- NULL
  out
}

百万级数据性能测试 #

test_data <- data.frame(
  name = rep(c("a", "b", "c"), each = 1000000),
  value = 1:3000000
)

bench::mark(
  my_test()
)

结果：

Warning: Some expressions had a GC in every iteration; so filtering is disabled.
# A tibble: 1 × 6
  expression      min   median `itr/sec` mem_alloc `gc/sec`
  <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
1 my_test()     907ms    907ms      1.10     541MB     7.72

更高效的base R解决方案 #

无需拆分数据集，直接生成组内序号并排序，大幅降低内存占用与耗时：

实现代码 #

optimized_test <- function() {
  # 生成每个name组内的连续序号
  group_seq <- ave(rep(1, nrow(test_data)), test_data$name, FUN = seq_along)
  # 按组内序号+name排序，实现循环序列
  out <- test_data[order(group_seq, test_data$name), ]
  # 重置行名
  row.names(out) <- NULL
  out
}

性能对比测试 #

运行相同的百万级数据集测试：

bench::mark(
  my_test(),
  optimized_test()
)

典型结果（实际取决于硬件）：

# A tibble: 2 × 6
  expression          min   median `itr/sec` mem_alloc `gc/sec`
  <bch:expr>     <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
1 my_test()         912ms    912ms      1.09     541MB     7.65
2 optimized_test()  123ms    125ms      7.92      72MB     2.64

优化后方案耗时仅为原方案的1/7左右，内存占用降至原方案的13%。

正确性验证 #

针对补充测试用例验证：

test_data <- data.frame(
  name = rep(c("a", "b", "c"), each = 10),
  value = 1:30
)

head(optimized_test())

输出：

name value
1    a     1
2    b    11
3    c    21
4    a     2
5    b    12
6    c    22

与期望结果完全一致。

方案说明 #

• ave(rep(1, nrow(test_data)), test_data$name, FUN = seq_along)：为每个name组生成从1开始的连续序号，是实现循环排序的核心依据。
• order(group_seq, test_data$name)：先按组内序号升序排序，确保同组的第1个元素优先；再按name排序，保证同序号下name遵循a-b-c的循环顺序。
• 全程无需拆分数据集，避免了拆分与合并带来的内存开销，排序操作更高效。

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