Skip to main content
 首页 » 编程设计

python之为什么我的 numpy.random.choice 实现速度更快

2025年05月04日16lyj

我想实现 numpy.random.choice (除了它的 replace 参数)看看它是如何工作的。

This is what I came up with :

from random import uniform 
from math import fsum 
 
def select(array, total_count, probability): 
    probability_accumulative = [] 
    last_element = 0 
    for i in range(len(probability)): 
        probability_accumulative.append(last_element + probability[i]) 
        last_element = probability_accumulative[i] 
 
    result = [] 
 
    if(len(array) != len(probability)): 
        raise ValueError("array and probability must have the same size.") 
    elif(fsum(probability) != 1.0): 
        raise ValueError("probabilities do not sum to 1.") 
    else: 
        for i in range(total_count): 
            rand = uniform(0, 1) 
            for j in range(len(probability_accumulative)): 
                if(rand < probability_accumulative[j]): 
                    result.append(array[j]) 
                    break 
 
    return result 

它似乎工作得很好,所以我决定编写另一个脚本来检查与 numpy.random.choice 相比,我的实现速度有多慢。

This is the test script I wrote for it :

from random_selection import select 
from collections import Counter 
from numpy.random import choice 
from time import time 
 
def test(array, total_count, probability, method): 
    methods = { 
        "numpy.random.choice": choice(array, total_count, p=probability), 
        "random_selection.select": select(array, total_count, probability) 
    } 
 
    if(method in methods): 
        probability_dict = {} 
        rand_items = methods[method] 
        items_counter = Counter(rand_items) 
 
        for item, count in items_counter.most_common(): 
            probability_dict[item] = f"{100 * count / total_count:.1f}%" 
        return probability_dict 
    else: 
        raise ValueError(f"Method {method} has not been defined.") 
 
 
def main(): 
    total_count = 1000000 
    array = ['a', 'b', 'c', 'd'] 
    probability = [0.7, 0.1, 0.1, 0.1] 
 
    print(f"array: {array}") 
    print(f"probability: {probability}") 
    print(f"size: {total_count}") 
 
    print() 
 
    print('random_selection.select: ') 
    start_time = time() 
    result = test(array, total_count, probability, 'random_selection.select') 
    end_time = time() 
    print(result) 
    print(f"{(end_time - start_time):.4f} s") 
 
    print() 
 
    print('numpy.random.choice: ') 
    start_time = time() 
    result = test(array, total_count, probability, 'numpy.random.choice') 
    end_time = time() 
    print(result) 
    print(f"{(end_time - start_time):.4f} s") 
 
 
if __name__ == "__main__": 
    main() 

我很惊讶地发现我的实现速度更快!

这是一百万数组大小的结果:

array: ['a', 'b', 'c', 'd'] 
probability: [0.7, 0.1, 0.1, 0.1] 
size: 1000000 
 
random_selection.select: 
{'a': '70.0%', 'c': '10.0%', 'd': '10.0%', 'b': '10.0%'} 
2.5119 s 
 
numpy.random.choice: 
{'a': '70.0%', 'b': '10.0%', 'd': '10.0%', 'c': '10.0%'} 
3.1098 s 

如果我将大小增加到 1000 万,差异会变得更加明显:

array: ['a', 'b', 'c', 'd'] 
probability: [0.7, 0.1, 0.1, 0.1] 
size: 10000000 
 
random_selection.select: 
{'a': '70.0%', 'b': '10.0%', 'd': '10.0%', 'c': '10.0%'} 
25.6174 s 
 
numpy.random.choice: 
{'a': '70.0%', 'b': '10.0%', 'c': '10.0%', 'd': '10.0%'} 
31.8087 s 

这是为什么?

请您参考如下方法:

您的测试代码没有达到您的预期。 测试 函数总是 调用您的两个随机选择函数两者。您的计时仅检测您的分析代码在与所请求函数对应的结果上的性能差异。

问题在于这些行:

methods = { 
    "numpy.random.choice": choice(array, total_count, p=probability), 
    "random_selection.select": select(array, total_count, probability) 
} 

这些无条件调用choice和select函数,并将返回值放入字典中。这几乎肯定不是您所期望的。您可能希望将 lambda 函数放入字典中,它会在调用时使用适当的参数调用所需的函数。