如何按Team分组、基于ID唯一值比例抽取指定行数随机样本

要解决小占比组抽0行的问题，同时保证总抽取行数不超过8，我们需要先给占比低于0.5%的团队预留1个名额，再将剩余名额按比例分配给其他团队，最后调整配额确保总数符合要求。以下是修改后的代码：

import pandas as pd

n_total = 8

# 1. 计算每个Team的唯一ID数量及占比
unique_counts = CaseLoad.groupby("Team")["ID"].nunique()
total_unique = unique_counts.sum()
ratios = unique_counts / total_unique

# 2. 标记占比低于0.5%的小团队，每个至少抽1行
small_team_threshold = 0.005
small_teams = ratios[ratios < small_team_threshold].index
num_small_teams = len(small_teams)

# 3. 初始化配额：小团队先分配1个名额
nums_selected = pd.Series(1, index=small_teams, dtype=int)

# 4. 处理剩余配额分配给非小团队
remaining_quota = n_total - num_small_teams
if remaining_quota > 0:
    # 筛选非小团队，计算其占比（基于自身总唯一ID数）
    large_teams = unique_counts.drop(small_teams)
    if not large_teams.empty:
        large_ratios = large_teams / large_teams.sum()
        large_quota = pd.Series.round(large_ratios * remaining_quota).astype(int)
        nums_selected = pd.concat([nums_selected, large_quota])
elif remaining_quota < 0:
    # 小团队数量超过总配额，随机选n_total个小团队各抽1行
    selected_small_teams = small_teams.sample(n=n_total, random_state=42)
    nums_selected = pd.Series(1, index=selected_small_teams, dtype=int)

# 5. 调整配额，确保总数量恰好为n_total，且小团队配额不低于1
current_total = nums_selected.sum()
if current_total != n_total:
    if current_total > n_total:
        # 超额时，从配额最多的非小团队里扣减（避免动小团队的1个名额）
        excess = current_total - n_total
        reducible_teams = nums_selected[(nums_selected > 1) | (~nums_selected.index.isin(small_teams))].sort_values(ascending=False).index
        for team in reducible_teams[:excess]:
            nums_selected[team] -= 1
    else:
        # 不足时，给配额最少的团队补1，优先小团队之外的
        deficit = n_total - current_total
        augmentable_teams = nums_selected.sort_values(ascending=True).index
        for team in augmentable_teams[:deficit]:
            nums_selected[team] += 1

# 6. 分组抽样，确保每个组抽取的数量不超过该组的总行数
Selection = main.groupby("Team", group_keys=False).apply(
    lambda x: x.sample(n=min(nums_selected.get(x.name, 0), len(x)), random_state=42, replace=False)
)

关键逻辑说明 #

• 小团队优先保障：先给占比<0.5%的团队分配1个名额，避免抽到0行。
• 剩余配额分配：剩余名额按非小团队的唯一ID占比分配，用四舍五入取整。
• 配额校准：如果分配后总数超过或不足8，通过调整配额最多/最少的团队（优先不改动小团队的1个名额），确保总数量恰好为8。
• 抽样边界处理：抽样时取配额和组内行数的较小值，避免出现sample函数因请求数量超过组内行数报错的情况。

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