按餐厅ID分层平衡布尔值的DataFrame最大子集提取问询

我来帮你实现这个按餐厅ID分层的均衡抽样需求！先明确你的核心诉求：

• 基于包含10000条样本、600个Restaurant_ID（含缺失）的数据集，其中Restaurant_Bool比例约84:16且各ID比例不同
• 按Restaurant_ID分层，每个ID下保留数量相等的0和1样本（取该ID下两类样本的最小数量）
• 剔除仅含0或仅含1的ID，保留餐厅内部相关性用于建模
• 支持生成多个不同的随机子集

先修正下你提供的原始模拟数据代码（原来的循环逻辑有问题，只会删除1个ID的样本），再分别给出R和Python的实现方案：

R版本 #

# 设置随机种子保证可复现
set.seed(123)
# 生成10000条样本，600个Restaurant_ID
x <- floor(runif(10000, 0, 600)) 
x <- sort(x)
# 生成偏态布尔值（84%为1，16%为0）
y <- sample(0:1, 10000, prob=c(.16,.84), replace=TRUE)
df <- data.frame(x, y)
colnames(df) <- c("Restaurant_ID","Restaurant_Bool")

# 删除10个随机ID的所有样本（修正原循环逻辑）
z <- floor(runif(10, 0, 600))
for (i in 1:10) {
  df <- df[!(df$Restaurant_ID == z[i]), ]
}

# 查看数据概况
summary(df)
nrow(df)

Python版本 #

import pandas as pd
import numpy as np

# 设置随机种子保证可复现
np.random.seed(123)
# 生成10000条样本，600个Restaurant_ID
x = np.floor(np.random.uniform(0, 600, 10000)).astype(int)
x = np.sort(x)
# 生成偏态布尔值（84%为1，16%为0）
y = np.random.choice([0,1], size=10000, p=[0.16, 0.84])
df = pd.DataFrame({'Restaurant_ID': x, 'Restaurant_Bool': y})

# 删除10个随机ID的所有样本
z = np.floor(np.random.uniform(0, 600, 10)).astype(int)
df = df[~df['Restaurant_ID'].isin(z)]

# 查看数据概况
print(df.describe())
print(f"样本总数：{len(df)}")

R实现方案（支持生成多个随机子集） #

library(dplyr)

# 定义均衡抽样函数
balanced_sample_by_id <- function(df, n_subsets = 1) {
  # 先过滤掉只有单类样本的ID
  valid_ids <- df %>%
    group_by(Restaurant_ID) %>%
    summarise(has_0 = any(Restaurant_Bool == 0),
              has_1 = any(Restaurant_Bool == 1)) %>%
    filter(has_0 & has_1) %>%
    pull(Restaurant_ID)
  
  filtered_df <- df %>% filter(Restaurant_ID %in% valid_ids)
  
  # 生成n_subsets个随机子集
  subsets <- list()
  for (i in 1:n_subsets) {
    subset_df <- filtered_df %>%
      group_by(Restaurant_ID, Restaurant_Bool) %>%
      sample_n(size = min(n()), replace = FALSE) %>% # 每个ID下取两类样本的最小数量
      ungroup()
    subsets[[i]] <- subset_df
  }
  
  return(subsets)
}

# 使用示例：生成3个不同的随机子集
set.seed(456)
my_subsets <- balanced_sample_by_id(df, n_subsets = 3)

# 查看第一个子集的概况
summary(my_subsets[[1]])
# 检查第一个子集每个ID的0/1数量是否相等
my_subsets[[1]] %>%
  group_by(Restaurant_ID, Restaurant_Bool) %>%
  count() %>%
  spread(Restaurant_Bool, n)

代码说明： #

• 首先筛选出同时包含0和1样本的Restaurant_ID，剔除单类ID
• 对每个有效ID，随机抽取该ID下0和1样本数量的最小值（保证两类样本数量相等）
• 通过循环生成指定数量的随机子集，每次抽样都是随机的，所以子集内容不同

Python实现方案（支持生成多个随机子集） #

def balanced_sample_by_id(df, n_subsets=1):
    # 先过滤掉只有单类样本的ID
    # 计算每个ID是否同时包含0和1
    id_stats = df.groupby('Restaurant_ID')['Restaurant_Bool'].agg(['nunique', 'count'])
    valid_ids = id_stats[id_stats['nunique'] == 2].index.tolist()
    filtered_df = df[df['Restaurant_ID'].isin(valid_ids)]
    
    subsets = []
    for _ in range(n_subsets):
        # 对每个ID，计算0和1的数量，取最小值作为抽样数
        sample_df = filtered_df.groupby('Restaurant_ID').apply(
            lambda group: group.groupby('Restaurant_Bool').apply(
                lambda g: g.sample(n=min(group['Restaurant_Bool'].value_counts()), replace=False)
            ).reset_index(drop=True)
        ).reset_index(drop=True)
        subsets.append(sample_df)
    
    return subsets

# 使用示例：生成3个不同的随机子集
np.random.seed(456)
my_subsets = balanced_sample_by_id(df, n_subsets=3)

# 查看第一个子集的概况
print(my_subsets[0].describe())
# 检查第一个子集每个ID的0/1数量是否相等
count_check = my_subsets[0].groupby(['Restaurant_ID', 'Restaurant_Bool']).size().unstack()
print(count_check)

代码说明： #

• 利用pandas的groupby功能，先筛选出同时包含0和1的ID
• 嵌套groupby实现按ID分层，再按Restaurant_Bool抽样，取两类样本的最小数量
• 循环生成多个随机子集，每次抽样都会随机选择样本，确保子集多样性

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