np.random.permutation、np.random.choice 的时间表现

Question

在我的纯 python 图论库中，我遇到了一个相对于可比较的 MATLAB 代码的时间性能非常差的函数，因此我尝试分析该函数中的一些操作。

我追踪到如下结果

In [27]: timeit.timeit( 'permutation(138)[:4]', setup='from numpy.random import permutation', number=1000000)
Out[27]: 27.659916877746582

将其与 MATLAB 中的性能进行比较

>> tic; for i=1:1000000; randperm(138,4); end; toc
Elapsed time is 4.593305 seconds.

通过将其更改为 np.random.choice 而不是我最初写的 np.random.permutation，我能够显着提高性能。

In [42]: timeit.timeit( 'choice(138, 4)', setup='from numpy.random import choice', number=1000000)
Out[42]: 18.9618501663208

但它仍然没有接近 matlab 的性能。

是否有另一种方法可以在纯 python 中获得这种行为，时间性能接近 MATLAB 时间性能？

Answer 1

基于this solution，它展示了如何使用基于argsort/argpartition 的技巧来模拟np.random.choice(..., replace=False) 的行为，您可以使用np.argpartition 重新创建MATLAB 的randperm(138,4)，即NumPy 的np.random.choice(138,4, replace=False)如：

np.random.rand(138).argpartition(range(4))[:4]

或者像这样使用np.argsort -

np.random.rand(138).argsort()[:4]

让我们对这两个版本进行计时，以便与 MATLAB 版本进行性能比较。

在 MATLAB 上 -

>> tic; for i=1:1000000; randperm(138,4); end; toc
Elapsed time is 1.058177 seconds.

在 NumPy 上使用 np.argpartition -

In [361]: timeit.timeit( 'np.random.rand(138).argpartition(range(4))[:4]', setup='import numpy as np', number=1000000)
Out[361]: 9.063489798831142

在 NumPy 上使用 np.argsort -

In [362]: timeit.timeit( 'np.random.rand(138).argsort()[:4]', setup='import numpy as np', number=1000000)
Out[362]: 5.74625801707225

最初建议使用 NumPy -

In [363]: timeit.timeit( 'choice(138, 4)', setup='from numpy.random import choice', number=1000000)
Out[363]: 6.793723535243771

似乎可以使用np.argsort 来稍微提高性能。

Answer 2

这对您来说需要多长时间？我估计 1-2 秒。

def four():
    k = np.random.randint(138**4)
    a = k % 138
    b = k // 138 % 138
    c = k // 138**2 % 138
    d = k // 138**3 % 138
    return (a, b, c, d) if a != b and a != c and a != d and b != c and b != d and c != d else four()

更新 1： 起初我使用 random.randrange，但 np.random.randint 使整个过程快了两倍。

更新 2：由于 NumPy 的随机函数似乎要快得多，所以我尝试了这个，它是另一个因素，速度快了 ~1.33：

>>> def four():
        a = randint(138)
        b = randint(138)
        c = randint(138)
        d = randint(138)
        return (a, b, c, d) if a != b and a != c and a != d and b != c and b != d and c != d else four()

>>> import timeit
>>> from numpy.random import randint
>>> timeit.timeit(lambda: four(), number=1000000)
2.3742770821572776

这比原来快了大约 22 倍：

>>> timeit.timeit('permutation(138)[:4]', setup='from numpy.random import permutation', number=1000000)
51.80568455893672

（字符串与lambda 没有明显区别）