为什么 random.sample 比 numpy 的 random.choice 快？

Question

我需要一种无需替换某个数组a 的采样方法。我尝试了两种方法（参见下面的 MCVE），使用 random.sample() 和 np.random.choice。

我认为numpy 函数会更快，但事实并非如此。在我的测试中，random.sample 比 np.random.choice 快约 15%。

这是正确的，还是我在下面的示例中做错了什么？如果这是正确的，为什么？

import numpy as np
import random
import time
from contextlib import contextmanager


@contextmanager
def timeblock(label):
    start = time.clock()
    try:
        yield
    finally:
        end = time.clock()
        print ('{} elapsed: {}'.format(label, end - start))


def f1(a, n_sample):
    return random.sample(range(len(a)), n_sample)


def f2(a, n_sample):
    return np.random.choice(len(a), n_sample, replace=False)


# Generate random array
a = np.random.uniform(1., 100., 10000)
# Number of samples' indexes to randomly take from a
n_sample = 100
# Number of times to repeat functions f1 and f2
N = 100000

with timeblock("random.sample"):
    for _ in range(N):
        f1(a, n_sample)

with timeblock("np.random.choice"):
    for _ in range(N):
        f2(a, n_sample)

Answer 1

TL;DR 从 numpy v1.17.0 开始，建议使用 numpy.random.default_rng() 对象而不是 numpy.random。供选择：

import numpy as np

rng = np.random.default_rng()    # you can pass seed
rng.choice(...)    # interface is the same

除了 v1.17 中引入的随机 API 的其他更改之外，这个新版本的选择现在更加智能，并且在大多数情况下应该是最快的。 为了向后兼容，旧版本保持不变！

---

正如 cmets 中所提到的，numpy 中存在一个长期存在的问题，即 np.random.choice 与 python 标准库中的 random.sample 相比，k << n 的实现无效。

问题是np.random.choice(arr, size=k, replace=False) 被实现为permutation(arr)[:k]。在大数组和小k的情况下，计算整个数组排列是浪费时间和内存。标准 python 的 random.sample 以更直接的方式工作 - 它只是迭代采样而不进行替换，或者跟踪已经采样的内容或要采样的内容。

在 v1.17.0 中，numpy 引入了对 numpy.random 包（docs、what's new、performance）的返工和改进。我强烈建议至少看看第一个链接。请注意，正如那里所说，为了向后兼容，旧的 numpy.random API 保持不变 - 它继续使用旧的实现。

所以使用随机 API 的新推荐方法是使用 numpy.random.default_rng() object 而不是 numpy.random。请注意，它是一个对象，它也接受可选的种子参数，因此您可以以方便的方式传递它。默认情况下，它还使用不同的生成器，平均速度更快（有关详细信息，请参阅上面的性能链接）。

关于您的情况，您现在可能想使用np.random.default_rng().choice(...)。除了速度更快之外，由于改进了随机生成器，choice 本身也变得更加智能。现在它只对足够大的数组（>10000 个元素）和相对大的 k（>1/50 的大小）使用整个数组排列。否则它使用 Floyd 的采样算法（short description，numpy implementation）。

---

这是我笔记本电脑上的性能比较：

来自 10000 个元素 x 10000 次的数组的 100 个样本：

random.sample elapsed: 0.8711776689742692
np.random.choice elapsed: 1.9704092079773545
np.random.default_rng().choice elapsed: 0.818919860990718

来自 10000 个元素 x 10000 次数组的 1000 个样本：

random.sample elapsed: 8.785315042012371
np.random.choice elapsed: 1.9777243090211414
np.random.default_rng().choice elapsed: 1.05490942299366

来自 10000 个元素 x 10000 次数组的 10000 个样本：

random.sample elapsed: 80.15063399000792
np.random.choice elapsed: 2.0218082449864596
np.random.default_rng().choice elapsed: 2.8596064270241186

以及我使用的代码：

import numpy as np
import random
from timeit import default_timer as timer
from contextlib import contextmanager


@contextmanager
def timeblock(label):
    start = timer()
    try:
        yield
    finally:
        end = timer()
        print ('{} elapsed: {}'.format(label, end - start))


def f1(a, n_sample):
    return random.sample(range(len(a)), n_sample)


def f2(a, n_sample):
    return np.random.choice(len(a), n_sample, replace=False)


def f3(a, n_sample):
    return np.random.default_rng().choice(len(a), n_sample, replace=False)


# Generate random array
a = np.random.uniform(1., 100., 10000)
# Number of samples' indexes to randomly take from a
n_sample = 100
# Number of times to repeat tested functions
N = 100000

print(f'{N} times {n_sample} samples')
with timeblock("random.sample"):
    for _ in range(N):
        f1(a, n_sample)

with timeblock("np.random.choice"):
    for _ in range(N):
        f2(a, n_sample)

with timeblock("np.random.default_rng().choice"):
    for _ in range(N):
        f3(a, n_sample)