如何用Pandas DataFrame实现多条件分层比例随机抽样

我来帮你梳理下需求并给出正确的实现方案，你的核心需求是两层分层抽样：先从每个Taluka筛选出数据量足够的村庄并随机选10个，再在每个选中村庄内按Group的占比抽取固定数量的样本。咱们一步步来拆解实现：

核心思路拆解 #

1. 筛选合格村庄：先过滤掉每个Taluka内数据量不足5条的村庄；
2. 抽取目标村庄：在每个Taluka的合格村庄中随机选10个（如果合格村庄不足10个则全选）；
3. 按Group比例抽样本：对每个选中的村庄，根据Group的记录占比计算各Group需抽取的数量，最终每个村庄抽取5条数据（自动调整比例取整后的总和，确保总数为5）。

完整实现代码 #

import pandas as pd
import numpy as np

# 读取数据
df = pd.read_excel("/home/desktop/Music/Data-Balaghat.xlsx")

# --------------------------
# 步骤1：筛选出每个Taluka内数据量≥5的村庄
# --------------------------
# 计算每个Taluka-Village组合的记录数
village_counts = df.groupby(['Taluka', 'Village'])['NDVI'].count().reset_index(name='count')
# 筛选出记录数≥5的村庄列表
qualified_villages = village_counts[village_counts['count'] >= 5][['Taluka', 'Village']]
# 合并回原数据，只保留合格村庄的记录
df_qualified = pd.merge(df, qualified_villages, on=['Taluka', 'Village'], how='inner')

# --------------------------
# 步骤2：在每个Taluka中随机选取10个合格村庄
# --------------------------
def select_villages_per_taluka(group):
    # 如果该Taluka的合格村庄数≥10，随机选10个；否则全选
    n_select = min(10, len(group['Village'].unique()))
    selected_villages = group['Village'].unique()
    np.random.shuffle(selected_villages)
    return group[group['Village'].isin(selected_villages[:n_select])]

# 按Taluka分组，执行村庄选取
df_selected_villages = df_qualified.groupby('Taluka').apply(select_villages_per_taluka).reset_index(drop=True)

# --------------------------
# 步骤3：在每个选中村庄内按Group比例抽取5条数据
# --------------------------
def sample_by_group_proportion(group):
    total_samples = 5
    # 计算该村庄内每个Group的记录占比
    group_counts = group['Group'].value_counts(normalize=True)
    # 计算初始抽样数量（四舍五入）
    sample_nums = np.round(group_counts * total_samples).astype(int)
    
    # 调整抽样数量，确保总和为5（处理四舍五入后的偏差）
    diff = total_samples - sample_nums.sum()
    if diff != 0:
        # 找到占比最大的Group，调整其抽样数量（diff为正就加，负就减）
        max_prop_group = group_counts.idxmax()
        sample_nums[max_prop_group] += diff
    
    # 按计算出的数量对每个Group进行抽样
    sampled = []
    for group_name, num in sample_nums.items():
        if num > 0:
            group_subset = group[group['Group'] == group_name]
            # 随机抽取num条，若该Group记录数不足num则全选
            sampled.append(group_subset.sample(n=min(num, len(group_subset)), random_state=42))
    return pd.concat(sampled)

# 按Taluka-Village分组，执行按Group比例抽样
final_sample = df_selected_villages.groupby(['Taluka', 'Village']).apply(sample_by_group_proportion).reset_index(drop=True)

# 查看最终抽样结果
print(final_sample[['Taluka', 'Village', 'Group', 'NDVI']].head())

代码关键点说明 #

• 合格村庄筛选：通过groupby(['Taluka', 'Village'])统计每个村庄的记录数，确保只保留数据量足够的村庄；
• 随机选村庄：用np.random.shuffle打乱村庄列表后选取前10个，保证随机性；加入了对村庄数量不足10的处理，避免报错；
• Group比例抽样：先计算占比得到初始抽样数，再自动调整总和到5，解决了四舍五入后总数偏差的问题；抽样时如果某个Group的记录数不足所需数量，直接全选该组所有数据，避免逻辑错误；
• 可复现性：在抽样时加入random_state=42，如果需要不同的随机结果可以去掉或修改该参数。

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