如何高效在R中按规则分配列值至其他列（避免循环）

方案1：dplyr实现（代码简洁，可读性高） #

library(dplyr)

df_processed <- df %>%
  group_by(key) %>%
  mutate(
    sum_y = sum(y),
    sum_z = sum(z),
    need_process = first(x) < sum_y & sum_z > first(x)
  ) %>%
  arrange(desc(z), .by_group = TRUE) %>%
  mutate(
    cum_z = cumsum(z),
    remaining_x = first(x) - lag(cum_z, default = 0),
    new_y = ifelse(need_process, pmin(y, pmax(remaining_x, 0)), y),
    # 修正最后一行，确保分配总和严格等于x
    new_y = ifelse(need_process & row_number() == n(),
                   pmax(first(x) - sum(new_y[-n()]), 0),
                   new_y)
  ) %>%
  select(key, x, y = new_y, z) %>%
  ungroup()

方案2：data.table实现（性能最优，适合10万+行） #

library(data.table)

setDT(df)

df_processed <- df[, 
  c("sum_y", "sum_z", "need_process") := .(sum(y), sum(z), first(x) < sum(y) & sum(z) > first(x)),
  by = key
][order(-z), 
  `:=`(cum_z = cumsum(z), remaining_x = first(x) - shift(cum_z, fill = 0)),
  by = key
][,
  new_y := fifelse(need_process, pmin(y, pmax(remaining_x, 0)), y),
  by = key
][need_process == TRUE,
  new_y := fifelse(.I == .I[.N], pmax(first(x) - sum(new_y[-.N]), 0), new_y),
  by = key
][, .(key, x, y = new_y, z)]

关键逻辑说明 #

1. 分组判断条件：先计算每组的sum_y和sum_z，标记需要处理的组
2. 排序优先级：对需处理的组按z降序排列，保证z值大的行优先分配
3. 分配规则：计算每行可分配的剩余x，取原y值和剩余x的最小值，避免分配超出原y的上限
4. 总和修正：最后一行手动调整，确保分配的y值总和严格等于组内x，消除累积计算误差

注意事项 #

• 假设每个key组的x值完全相同（符合问题描述），用first(x)提取组内x值
• 两种方案均采用向量化操作，无for循环，适合大规模数据处理
• data.table在处理10万+行数据时，分组、排序、赋值的性能显著优于dplyr

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