高效随机采样实现方案问询：多列表按概率采样优化

现有四个列表A、B、C、D，其中列表A包含n_a个宽度为1的元素，列表B包含n_b个宽度为2的元素，列表C包含n_c个宽度为3的元素，列表D包含n_d个宽度为4的元素。目标是根据用户定义的概率分布（如{1: 0.25, 2: 0.25, 3: 0.5, 4: 0.0}），从这四个列表中随机采样元素构建数组。例如，使用该分布时，输出数组中应约有25%来自列表A的1宽度元素、25%来自列表B的2宽度元素、50%来自列表C的3宽度元素，且不包含4宽度元素。

当前采用如下代码实现随机采样：

import numpy as np
import random

A = [1, 2, 3]
B = ['aa', 'bb', 'cc']
C = [111, 222, 333]
D = ['aaaa', 'bbbbb', 'cccc']

num_of_rows = 4  # 最终网格的行数
num_of_coulmns = 5  # 最终网格的列数

# 定义概率
p_A = 0.25
p_B = 0.25
p_C = 0.5
p_D = 0.0

probabilities = [p_A, p_B, p_C, p_D]

# 构建随机采样数组
result_array = []

for _ in range(num_of_rows):
    # 根据概率选择列表
    chosen_list = random.choices([A, B, C, D], weights=probabilities, k=1)[0]
    # 从选中列表中随机选元素
    element = random.choice(chosen_list)
    result_array.append(element)

但当网格行数和列数较大（如1000）时，该方法效率过低，请问是否有更优的实现方式？

核心思路是用numpy向量化操作替代Python循环，numpy底层基于C实现，批量处理的速度远快于纯Python循环。具体步骤如下：

• 一次性生成所有行对应的列表类型索引，避免循环调用随机选择函数
• 统计每种列表需要采样的数量，批量完成元素采样
• 最后打乱采样结果，保证随机性与原逻辑一致

优化后的代码：

import numpy as np

A = [1, 2, 3]
B = ['aa', 'bb', 'cc']
C = [111, 222, 333]
D = ['aaaa', 'bbbbb', 'cccc']

num_of_rows = 1000
probabilities = [0.25, 0.25, 0.5, 0.0]
lists = [A, B, C, D]

# 1. 批量生成所有行对应的列表索引（0对应A，1对应B，2对应C，3对应D）
list_indices = np.random.choice([0, 1, 2, 3], size=num_of_rows, p=probabilities)

# 2. 统计每个列表需要采样的元素数量
counts = np.bincount(list_indices, minlength=4)

# 3. 批量采样每个列表的元素
samples = []
for idx, cnt in enumerate(counts):
    if cnt == 0:
        continue
    # 从对应列表中随机选取cnt个元素（允许重复采样）
    sampled_elements = np.random.choice(lists[idx], size=cnt, replace=True)
    samples.extend(sampled_elements)

# 4. 打乱采样结果，保证每行的随机性
np.random.shuffle(samples)

# 转换为Python列表（如果需要）
result_array = samples.tolist()

优化效果说明 #

• 避免了原代码中逐行循环的开销，将多次随机调用合并为少数几次批量操作
• numpy的随机采样函数效率远高于Python标准库的random模块，尤其在数据量较大时优势明显
• 最终结果的概率分布与原逻辑完全一致，同时保证了元素选择的随机性

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