如何在Python中快速从超大数组随机采样指定数量元素？

当从超1000万个点中采样少量样本（如8192个）时，np.random.choice(np.arange(len(points)))会生成巨大的索引数组，既占用内存又拖慢执行速度。以下是几种高效的替代方案：

1. 优化原生numpy调用，避免生成大数组 #

np.random.choice支持直接传入整数作为总体规模，无需手动生成np.arange索引数组，这能大幅减少内存开销和初始化时间：

selected_point_idxs = np.random.choice(len(points), npoints, replace=False)

2. 使用numpy新随机数生成器（推荐） #

numpy 1.17+引入的np.random.Generator采用了更高效的采样算法，尤其适合采样量远小于总体规模的无放回采样场景，性能显著优于旧版API：

import numpy as np

rng = np.random.default_rng()  # 初始化随机数生成器
selected_point_idxs = rng.choice(len(points), npoints, replace=False)

3. 利用Python标准库random.sample #

random.sample可直接对惰性的range对象采样，无需生成全量索引，在部分场景下速度更快，最后转换为numpy数组即可：

import random
import numpy as np

selected_point_idxs = np.array(random.sample(range(len(points)), npoints))

4. 极端场景：手动实现蓄水池抽样 #

如果数据无法全部加载到内存（如流数据场景），可使用蓄水池抽样算法，仅需遍历一次数据即可完成采样，时间复杂度为O(N)：

import random
import numpy as np

def reservoir_sampling(data_size, sample_size):
    # 初始化样本池
    sample = list(range(sample_size))
    for i in range(sample_size, data_size):
        # 生成0到i的随机数，替换样本池中的元素
        rand_idx = random.randint(0, i)
        if rand_idx < sample_size:
            sample[rand_idx] = i
    return np.array(sample)

selected_point_idxs = reservoir_sampling(len(points), npoints)

方案选择建议 #

• 优先尝试方案2，新生成器的算法经过针对性优化，在绝大多数大样本小采样场景下表现最优；
• 方案1是最简便的改法，无需更换API，仅需删除多余的np.arange调用；
• 方案3适合numpy版本较低或纯Python环境；
• 方案4仅适用于数据无法全部加载到内存的流处理场景，你的数据已加载完成时，前三种方案足够满足需求。

内容的提问来源于stack exchange，提问作者Zhimin-J