给定N个变量取值范围生成K个无重复numpy随机数组的实现方法

需求明确 #

• 共N个独立变量可构成N维特征空间，当前场景N=13，每个变量对应预定义的整数取值列表，单列表长度约为20
• 每个变量的采样规则可自定义，支持均匀分布或正态分布两种模式
• 输出K个N维样本（兼容numpy数组格式），要求每个样本的第i个元素从第i个变量的取值列表中采样得到
• 仅禁止完整的N维样本完全重复，单个维度的取值允许重复出现

示例取值范围：

v1_range = [1, 2, 3, 4, 5, 6]
v2_range = [10, 20, 30, 40, 50, 60]
v3_range = [100, 200, 300, 400, 500]
v4_range = [15, 16, 17, 18]

符合要求的输出示例（K=6）：

output = [
    [1, 20, 300, 15],
    [5, 10, 200, 15],
    [3, 60, 100, 18],
    [4, 30, 200, 16],
    [4, 40, 100, 17],
    [2, 30, 400, 15],
]

需要避免的重复样本示例：

res = [[  4,  40, 400,  18],
       [  4,  60, 400,  17],
       [  4,  40, 400,  18]] # 与第0个样本完全重复，属于非法输出

最优实现方案 #

思路说明 #

当前场景的总样本空间为20^13，远大于常规业务场景下的K值，因此不需要枚举全量组合再采样，采用「批量生成+哈希判重+增量补全」的方案性能最优：

1. 用numpy.random.choice实现单维度自定义分布采样，传入p参数即可指定每个取值的采样概率，适配均匀、正态分布需求
2. 用集合存储已生成样本的元组哈希，O(1)复杂度判重，效率远高于逐行比对
3. 优先批量生成略多于K的样本，去重后如果满足需求直接返回前K个，不足的话再增量生成补充，减少循环次数

代码实现 #

import numpy as np
from scipy.stats import norm

def generate_unique_samples(var_ranges: list[list], k: int, dist_types: list[str] = None) -> np.ndarray:
    """
    生成无重复的N维随机样本
    :param var_ranges: 每个变量的取值列表，长度为N
    :param k: 需要生成的样本数量
    :param dist_types: 每个变量的分布类型，可选'uniform'(均匀)或'norm'(正态)，默认全为均匀分布
    :return: 形状为(k, N)的numpy数组
    """
    n = len(var_ranges)
    if dist_types is None:
        dist_types = ['uniform'] * n
    
    # 预计算每个维度的采样权重
    var_weights = []
    for i in range(n):
        val_num = len(var_ranges[i])
        if dist_types[i] == 'uniform':
            # 均匀分布权重相等
            var_weights.append(np.ones(val_num) / val_num)
        elif dist_types[i] == 'norm':
            # 正态分布以取值列表中间位置为均值，标准差设为长度的1/3
            mean = (val_num - 1) / 2
            std = val_num / 3
            x = np.arange(val_num)
            pdf = norm.pdf(x, loc=mean, scale=std)
            var_weights.append(pdf / pdf.sum())
    
    existed = set()
    res = []
    # 首次批量生成1.1倍K的样本，减少后续补全次数
    batch_size = int(k * 1.1)
    
    while len(res) < k:
        # 批量生成样本
        batch = np.zeros((batch_size, n), dtype=int)
        for i in range(n):
            batch[:, i] = np.random.choice(var_ranges[i], size=batch_size, p=var_weights[i])
        # 去重
        for row in batch:
            row_tuple = tuple(row)
            if row_tuple not in existed:
                existed.add(row_tuple)
                res.append(row)
                if len(res) == k:
                    break
        # 调整下一批次的生成数量
        batch_size = max(1, k - len(res))
    
    return np.array(res)

使用示例 #

# 示例参数
var_ranges = [
    [1, 2, 3, 4, 5, 6],
    [10, 20, 30, 40, 50, 60],
    [100, 200, 300, 400, 500],
    [15, 16, 17, 18]
]
# 前2个变量均匀分布，后2个变量正态分布
dist_types = ['uniform', 'uniform', 'norm', 'norm']
# 生成6个无重复样本
samples = generate_unique_samples(var_ranges, k=6, dist_types=dist_types)
print(samples)

性能说明 #

在13维每维20个取值的场景下，只要K不超过千万量级，该方案的运行时间都在毫秒级，碰撞概率极低，内存占用仅和K的大小成正比，远优于全组合枚举采样的方案。

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