为什么我的numpy.random.choice实现更快？

debugcn 发表于 Dev

阿米尔·沙巴尼（Amir A. Shabani）

我想实现numpy.random.choice（除了replace它的参数之外）以了解它是如何工作的。

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。

这是我为此编写的测试脚本：

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

这是为什么？

布莱克金

您的测试代码无法实现您期望的功能。该test功能一直呼吁双方你的两个随机选择功能。您的时间安排只能检测出与所需功能相对应的结果在分析代码上的性能差异。

问题在于以下几行：

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

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

本文收集自互联网，转载请注明来源。

如有侵权，请联系[email protected] 删除。