大DataFrame按分组优化for循环：独立计算colC列需求

Problem Recap #

You have a large DataFrame df where you need to compute colC with this logic:

• If day == 1, colC = colA + colB
• If day != 1, colC = (colA + colB) - previous row's colC

The data is split into consecutive groups of day 1-4, and each group needs to compute colC independently. The original for loop is too slow for large datasets.

Why the Original For Loop Is Slow #

R's base for loops are inefficient for large datasets because they perform row-wise operations, which involve repeated memory access and modification. For a dataset with millions of rows, this adds up to significant overhead.

Solution 1: Using dplyr + purrr (Readable & Efficient) #

We'll first create a group identifier for each consecutive day 1-4 block, then use purrr::accumulate to handle the recursive calculation of colC within each group. This is vectorized and much faster than a global for loop.

library(dplyr)
library(purrr)

# Sample input data
df <- tibble(
  day = c(1,2,3,4,1,2),
  colA = c(4,9,8,9,5,2),
  colB = c(8,2,1,4,8,4)
)

df <- df %>%
  # Create a group ID: increment every time we hit day = 1
  mutate(group_id = cumsum(day == 1)) %>%
  # Group by the new ID to process each day 1-4 block independently
  group_by(group_id) %>%
  mutate(
    sum_ab = colA + colB,
    # Use accumulate to handle the recursive colC calculation
    colC = accumulate(sum_ab, .init = 0, ~ ifelse(.y == sum_ab[1], .y, .y - .x))[-1]
  ) %>%
  ungroup() %>%
  # Clean up helper columns
  select(-group_id, -sum_ab)

# Check the result
print(df)

Solution 2: Using data.table (Best for Ultra-Large Datasets) #

If your dataset is extremely large (millions of rows), data.table offers even better performance due to its optimized grouping and memory management. We still group by consecutive day blocks, but use a tiny per-group for loop (only 3 iterations per group, since each group has 4 rows max) which is negligible in terms of overhead.

library(data.table)

# Convert to data.table
setDT(df)

# Create group ID
df[, group_id := cumsum(day == 1)]

# Compute colC per group
df[, colC := {
  sum_ab = colA + colB
  colC_vec = numeric(.N)
  colC_vec[1] = sum_ab[1]
  # Only loop 3 times per group (since groups are day 1-4)
  for(i in 2:.N) {
    colC_vec[i] = sum_ab[i] - colC_vec[i-1]
  }
  colC_vec
}, by = group_id]

# Remove helper column
df[, group_id := NULL]

# Check the result
print(df)

Expected Output #

Both solutions will produce your desired result:

# A tibble: 6 × 4
    day  colA  colB  colC
  <dbl> <dbl> <dbl> <dbl>
1     1     4     8    12
2     2     9     2    -1
3     3     8     1    10
4     4     9     4     3
5     1     5     8    13
6     2     2     4    -7

Performance Boost #

• The grouped approach reduces the total number of loop iterations drastically: instead of looping through every row in the entire dataset, you loop through at most 3 rows per group.
• dplyr and data.table use vectorized operations under the hood, which are far more efficient than base R's row-wise for loops.

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