基于count值随机生成指定行列数的Pandas DataFrame

问题背景 #

需要生成一个包含r行、列数由输入的group列决定的DataFrame，每个列的True值数量由对应的count列指定，要求随机分配这些True值。现有实现通过临时DataFrame配合explode方法完成，但存在以下疑问：

• Pandas是否有更简便的实现方法？
• 如何在expand()函数中实现随机分配？
• 是否可先创建含NaN的DataFrame再随机填充True值？

输入示例 #

假设总行数r=5（可自定义）。

现有实现代码 #

import pandas as pd
import numpy as np

# 输入数据
input_df = pd.DataFrame({'group': ['A', 'B', 'C', 'D'], 'count': [2, 0, 4, 1]})
r = 5

# 现有实现：临时DataFrame + explode
temp = input_df.assign(
    vals=lambda x: x['count'].apply(lambda n: [True]*n + [False]*(r-n))
).assign(
    vals=lambda x: x['vals'].apply(np.random.permutation)
).explode('vals')

result = temp.pivot(index=temp.groupby('group').cumcount(), columns='group', values='vals')

疑问解答与优化实现 #

1. Pandas是否有更简便的实现方法？ #

有，直接利用numpy的随机抽样结合DataFrame构造，代码更简洁：

import pandas as pd
import numpy as np

r = 5
input_df = pd.DataFrame({'group': ['A', 'B', 'C', 'D'], 'count': [2, 0, 4, 1]})

# 生成每列的布尔数组：随机选择count个位置为True，其余为False
data = {
    row['group']: np.random.choice([True, False], size=r, p=[row['count']/r, 1-row['count']/r])
    if row['count'] !=0 else [False]*r
    for _, row in input_df.iterrows()
}

result = pd.DataFrame(data)

这种方法直接通过字典推导生成每列数据，无需临时DF和explode，逻辑更直观。

2. 如何在expand()函数中实现随机分配？ #

如果要基于expand()实现随机分配，可以对每个group的行索引进行随机抽样，生成对应位置的True值：

import pandas as pd
import random

r = 5
input_df = pd.DataFrame({'group': ['A', 'B', 'C', 'D'], 'count': [2, 0, 4, 1]})

# 生成每个group的索引列表，随机选中count个设为True
def expand_row(row):
    indices = list(range(r))
    true_indices = random.sample(indices, row['count']) if row['count'] >0 else []
    return [True if i in true_indices else False for i in indices]

# 应用expand并构造DataFrame
expanded = input_df.apply(expand_row, axis=1)
result = pd.DataFrame(expanded.tolist(), index=input_df['group']).T

这里的expand_row函数针对每个group，随机抽取指定数量的行索引，生成布尔数组，再转置得到目标DF。

3. 是否可先创建含NaN的DataFrame再随机填充True值？ #

完全可以，这种方式灵活性更高，步骤如下：

import pandas as pd
import random

r = 5
input_df = pd.DataFrame({'group': ['A', 'B', 'C', 'D'], 'count': [2, 0, 4, 1]})

# 初始化全NaN的DataFrame
result = pd.DataFrame(index=range(r), columns=input_df['group'])

# 遍历每个group，随机填充指定数量的True
for _, row in input_df.iterrows():
    group = row['group']
    count = row['count']
    if count ==0:
        result[group] = False
        continue
    # 随机选择count个行索引
    true_rows = random.sample(list(result.index), count)
    result.loc[true_rows, group] = True
    # 剩余行设为False（可选，若保留NaN则跳过）
    result.loc[~result.index.isin(true_rows), group] = False

这种方法先构建空框架，再逐列填充，适合需要保留部分NaN的场景（只需去掉最后一行赋值即可）。

预期输出示例（r=5） #

A      B      C      D
0   True  False   True  False
1   True  False   True  False
2  False  False   True   True
3  False  False   True  False
4  False  False  False  False

注：由于随机分配，每次运行结果会有差异，但各列的True数量严格匹配输入的count值。

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