R语言：按ID分组缩减DataFrame，平衡组内不同rating行数

Hey there! Let's tackle this problem where we need to make sure each id group has an equal number of rows with rating=0 and rating=1. First, let's recap your dataset and what we're aiming for:

Original Dataset #

Your raw data looks like this (with code to recreate it):

id <- c("vp01","vp01","vp01","vp01", "vp02", "vp02","vp02","vp02","vp02", "vp02")
rating <- c("0","1","0","0","1","0","1", "0", "0", "0")
Au1 <- c("150.0","100.45","80.23","133.21","94.33","102.22", "83.45", "122.65", "115.41", "109.34")
df <- data.frame(id,rating,Au1)

Key Goal #

For each id:

• Count how many rows have rating=1 and rating=0
• Keep the smaller count of rows for both ratings (so they match)
• For example:
  • vp01 has 1 row with rating=1, so we keep 1 row with rating=0
  • vp02 has 2 rows with rating=1, so we keep 2 rows with rating=0

Approach 1: Using dplyr (Tidyverse Style) #

This method is clean and readable, perfect for data manipulation tasks. First, we'll fix the data types (since rating and Au1 are stored as characters right now) then balance the groups.

# Load the dplyr package (install it first if you haven't: install.packages("dplyr"))
library(dplyr)

# Convert columns to appropriate data types
df_clean <- df %>%
  mutate(
    rating = as.integer(rating),
    Au1 = as.numeric(Au1)
  )

# Balance the rows per id group
balanced_df <- df_clean %>%
  group_by(id) %>%
  # Calculate how many 1s, 0s, and the minimum of the two
  mutate(
    count_1 = sum(rating == 1),
    count_0 = sum(rating == 0),
    target_count = min(count_1, count_0)
  ) %>%
  # Group again by id AND rating to sample the target number of rows
  group_by(id, rating) %>%
  # Use slice_head to take the first N rows (matches your example)
  # Use slice_sample instead if you want random rows instead of first ones
  slice_head(n = first(target_count)) %>%
  # Remove helper columns and ungroup
  ungroup() %>%
  select(-count_1, -count_0, -target_count) %>%
  # Sort to match your expected output
  arrange(id, rating)

# View the result
balanced_df

Running this will give you exactly the output you wanted:

# A tibble: 6 × 3
  id    rating   Au1
  <chr>  <int> <dbl>
1 vp01       0  150 
2 vp01       1  100.
3 vp02       0  102.
4 vp02       0  123.
5 vp02       1   94.3
6 vp02       1   83.4

Approach 2: Using Base R #

If you prefer not to use external packages, here's a base R solution that does the same thing:

# Fix data types first
df$rating <- as.integer(df$rating)
df$Au1 <- as.numeric(df$Au1)

# Split the data by id and process each group
balanced_groups <- lapply(split(df, df$id), function(group) {
  # Count 1s and 0s in the group
  num_1 <- sum(group$rating == 1)
  num_0 <- sum(group$rating == 0)
  keep_num <- min(num_1, num_0)
  
  # Select the first N rows for each rating
  rows_1 <- group[group$rating == 1, ][1:keep_num, ]
  rows_0 <- group[group$rating == 0, ][1:keep_num, ]
  
  # Combine and return the balanced group
  rbind(rows_1, rows_0)
})

# Combine all groups back into a single data frame
balanced_df_base <- do.call(rbind, balanced_groups)
# Reset row names to be sequential
rownames(balanced_df_base) <- NULL
# Sort to match the expected output
balanced_df_base <- balanced_df_base[order(balanced_df_base$id, balanced_df_base$rating), ]

# View the result
balanced_df_base

This will produce the same balanced dataset as the dplyr method.

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