Polars DataFrame中列表的加权抽样实现问询

一、从每个组的列表中抽取固定数量元素 #

要从elements列的每个列表中抽取5个元素，直接使用Polars的list.sample()表达式即可，该方法会对每行的列表单独进行随机抽样：

import polars as pl

df = pl.DataFrame({
    "group": ["foo", "bar", "baz"],
    "elements": [
        pl.int_range(0, 100, eager=True),
        pl.int_range(200, 300, eager=True),
        pl.int_range(300, 400, eager=True)
    ],
    "weight": [0.1, 0.5, 0.4]
})

# 每个列表抽取5个元素，设置seed保证结果可复现
df_sampled_per_group = df.with_columns(
    pl.col("elements").list.sample(n=5, seed=42).alias("elements")
)

print(df_sampled_per_group)

输出结果类似：

┌───────┬───────────────────────┬────────┐
│ group ┆ elements              ┆ weight │
│ ---   ┆ ---                   ┆ ---    │
│ str   ┆ list[i64]             ┆ f64    │
╞═══════╪═══════════════════════╪════════╡
│ foo   ┆ [7, 42, 19, 74, 33]   ┆ 0.1    │
│ bar   ┆ [209, 277, 222, 291, 260] ┆ 0.5 │
│ baz   ┆ [300, 347, 312, 398, 369] ┆ 0.4 │
└───────┴───────────────────────┴────────┘

二、按权重从所有组中抽取指定总数的元素 #

要按照weight列的权重总共抽取1000个元素，步骤如下：

1. 计算每个组的抽样数量：总样本数 × 组权重，转换为整数
2. 对每个组的elements列表抽取对应数量的元素
3. 可选：将抽样结果展开为单行一个元素的格式

代码实现： #

sample_total = 1000

# 计算每个组需要抽取的元素数量
df_with_sample_count = df.with_columns(
    (pl.col("weight") * sample_total).round().cast(pl.Int32).alias("sample_count")
)

# 按计算出的数量对每个组的列表抽样
df_weighted_sampled = df_with_sample_count.with_columns(
    pl.col("elements").list.sample(n=pl.col("sample_count"), seed=42).alias("sampled_elements")
)

# 若需要将所有抽样元素展开为单行一个元素的格式
df_weighted_sampled_exploded = df_weighted_sampled.select("group", "sampled_elements").explode("sampled_elements")

# 查看结果行数（应为1000）
print(len(df_weighted_sampled_exploded))

说明： #

• 这里使用round()处理权重转整数的问题，若权重计算后的总和与目标总数有偏差，可根据需求调整（比如使用floor()/ceil()，或手动微调某组的抽样数量）
• 设置seed参数可以保证抽样结果的可复现性
• 展开后的数据集每行对应一个抽样元素，保留了原组的信息

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