如何为每个ID抽取1行样本且保证组间样本数量均等？

现有一数据集，每个ID对应多行数据，共约5000个ID，每个ID的行数为1至22行，每行分属不同group。需求为：为每个ID抽取1行样本，且最终样本中各组的样本数量均等。

示例数据集 #

简化示例含8个ID，每个ID对应1至4行数据：

id group
1  a
1  b
1  c
1  d
2  a
2  b
3  a
3  b
3  c
3  d
4  a
4  b
4  d
5  a
5  b
5  c
5  d
6  a
6  d
7  a
7  b
7  d
8  a
8  b
8  c
8  d

期望结果 #

因共8个ID、4个group，最终样本需每个group对应2个ID（如20个ID、4个group则每个group对应5个ID），且所有ID仅出现一次。示例输出如下：

id  group
1   d
2   b
3   a
4   d
5   c
6   a
7   b
8   c

Python解决方案 #

直接用权重抽样无法保证严格的组间均衡，这里通过贪心分配+随机调整实现：

步骤1：构建示例数据 #

import pandas as pd
import numpy as np

d = {'id': [1,1,1,1,
            2,2,
            3,3,3,3,
            4,4,4,
            5,5,5,5,
            6,6,
            7,7,7,
            8,8,8,8], 
     'group': ['a','b','c','d',
            'a','b',
            'a','b','c','d',
            'a','b','d',
            'a','b','c','d',
            'a','d',
            'a','b','d',
            'a','b','c','d']}
df = pd.DataFrame(data=d)

步骤2：实现均衡抽样 #

# 计算每个group需要分配的ID数量
total_ids = df['id'].nunique()
groups = df['group'].unique()
target_per_group = total_ids // len(groups)

# 按ID分组，记录每个ID可选的group
id_groups = df.groupby('id')['group'].apply(list).reset_index(name='available_groups')

# 初始化结果字典和组计数
result = {}
group_count = {g:0 for g in groups}

# 先处理只有单个可选group的ID，避免无法分配
for idx, row in id_groups.iterrows():
    if len(row['available_groups']) == 1:
        g = row['available_groups'][0]
        result[row['id']] = g
        group_count[g] += 1

# 处理剩余ID，优先分配给未达目标的组
remaining_ids = id_groups[~id_groups['id'].isin(result.keys())]
for idx, row in remaining_ids.iterrows():
    # 筛选当前ID可选且未达目标的group
    eligible_groups = [g for g in row['available_groups'] if group_count[g] < target_per_group]
    # 随机选一个符合条件的group
    selected_g = np.random.choice(eligible_groups)
    result[row['id']] = selected_g
    group_count[selected_g] += 1

# 转换为DataFrame并排序
final_df = pd.DataFrame.from_dict(result, orient='index', columns=['group']).reset_index().rename(columns={'index':'id'})
final_df = final_df.sort_values('id').reset_index(drop=True)
print(final_df)

说明 #

• 先锁定只能选单个group的ID，避免这类ID无法分配导致组计数失衡
• 剩余ID优先分配给未达目标数量的组，保证最终各组样本数严格均等
• 随机选择符合条件的group，保留抽样随机性

R语言解决方案 #

采用相同的贪心分配逻辑：

步骤1：构建示例数据 #

library(dplyr)
library(tibble)

d <- list(
  id = c(1,1,1,1,
         2,2,
         3,3,3,3,
         4,4,4,
         5,5,5,5,
         6,6,
         7,7,7,
         8,8,8,8),
  group = c('a','b','c','d',
            'a','b',
            'a','b','c','d',
            'a','b','d',
            'a','b','c','d',
            'a','d',
            'a','b','d',
            'a','b','c','d')
)
df <- as_tibble(d)

步骤2：实现均衡抽样 #

# 计算每个group的目标数量
total_ids <- length(unique(df$id))
groups <- unique(df$group)
target_per_group <- total_ids %/% length(groups)

# 按ID整理可选group
id_groups <- df %>%
  group_by(id) %>%
  summarise(available_groups = list(group), .groups = 'drop')

# 初始化结果和计数
result <- list()
group_count <- setNames(rep(0, length(groups)), groups)

# 处理只有单个可选group的ID
single_group_ids <- id_groups %>% filter(lengths(available_groups) == 1)
for (i in 1:nrow(single_group_ids)) {
  id_val <- single_group_ids$id[i]
  g_val <- single_group_ids$available_groups[[i]][1]
  result[[as.character(id_val)]] <- g_val
  group_count[g_val] <- group_count[g_val] + 1
}

# 处理剩余ID
remaining_ids <- id_groups %>% filter(!id %in% as.numeric(names(result)))
for (i in 1:nrow(remaining_ids)) {
  id_val <- remaining_ids$id[i]
  eligible_groups <- remaining_ids$available_groups[[i]][group_count[remaining_ids$available_groups[[i]]] < target_per_group]
  selected_g <- sample(eligible_groups, 1)
  result[[as.character(id_val)]] <- selected_g
  group_count[selected_g] <- group_count[selected_g] + 1
}

# 转换为数据框并排序
final_df <- enframe(result, name = 'id', value = 'group') %>%
  mutate(id = as.numeric(id)) %>%
  arrange(id)
print(final_df)

说明 #

• 逻辑与Python版本一致，先处理无选择空间的ID，再对剩余ID做均衡分配
• 使用sample()实现随机选择，保证抽样的随机性

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