Python中无需转列表从大型字典随机抽样的方法

问题核心 #

• 直接调用random.sample(my_dict, k)会报错，因为字典不属于序列类型
• 使用random.sample(my_dict.keys(), k)会触发弃用警告，Python 3.9+已不再支持从集合类对象采样
• 把字典键转成列表会带来巨大内存开销，多次调用random.choice既会出现重复采样，又因字典不是序列直接抛出异常

最优方案：蓄水池抽样算法 #

针对超大字典，蓄水池抽样可以在仅遍历一次字典的前提下，完成无放回随机抽取k个样本的需求，无需将所有键/值加载到内存，彻底规避内存开销问题。

代码实现 #

import random

def random_dict_sample(d, k):
    """从字典d中无放回随机抽取k个键（如需值可修改返回逻辑）"""
    if k <= 0:
        return []
    if k > len(d):
        raise ValueError("Sample size cannot exceed the number of items in the dictionary")
    
    sample = []
    # 遍历字典的键（如需直接抽值，可改为遍历d.values()）
    for idx, key in enumerate(d):
        if idx < k:
            # 前k个元素直接加入样本池
            sample.append(key)
        else:
            # 生成随机位置，替换样本池中的元素
            replace_pos = random.randint(0, idx)
            if replace_pos < k:
                sample[replace_pos] = key
    return sample

# 示例用法
my_large_dict = {f"key_{i}": f"value_{i}" for i in range(1000000)}
# 抽取5个随机键
random_keys = random_dict_sample(my_large_dict, 5)
# 获取对应的值
random_values = [my_large_dict[key] for key in random_keys]

方案优势 #

• 内存开销固定为O(k)，仅维护一个长度为k的样本列表，与字典总大小无关
• 仅需遍历一次字典，时间复杂度为O(n)（n为字典元素数量）
• 可灵活调整为抽取键或值，只需修改遍历对象和样本存储内容

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