如何使用像numpy.random.choice这样的Polars基于概率对值进行采样？

Question

我正在使用 Python 中的 Polars 库，并希望根据相关概率从列表中采样值，类似于 numpy.random.choice 的工作方式。这就是我想要实现的目标：

在 numpy 中，我可以做这样的事情：

import numpy as np

# Possible outcomes
values = [0, 1, 2, 3, 4]

# Associated probabilities
probabilities = [0.15, 0.30, 0.25, 0.20, 0.10]

# Sample a value based on probabilities
sampled_value = np.random.choice(values, p=probabilities)
print(sampled_value)

这会从值中返回一个随机值，根据概率中的概率选择。

但是，在我使用 Polars 的用例中，概率位于 DataFrame 的单独列中，如下所示：

import polars as pl

df = pl.DataFrame({
    "value_0": [0.15],
    "value_1": [0.30],
    "value_2": [0.25],
    "value_3": [0.20],
    "value_4": [0.10]
})

有没有办法直接在 Polars DataFrame 中根据这些概率进行采样，或者我需要先将 DataFrame 转换为不同的格式？理想情况下，如果可能的话，我希望以矢量化方式执行此操作。

任何指导或示例将不胜感激。谢谢！

Answer 1

假设您有一行数据并想要对列名称进行采样，您可以使用

numpy.choice

，如您的示例所示：

np.random.choice(df.columns, p=df.row(0))

输出示例：

'value_3'