понедельник, 11 апреля 2011 г.

Взвешенный выбор в Python

Очень часто встречаю людей, которые зачем-то пытаются сделать взвешенный выбор в питоне, создавая огромных размеров массивы с повторяющимися значениями и потом выбирая из них результат с помощью random.choice. Вариант, конечно, логичный и простой в реализации, но... как же он тормозит и сколько памяти расходует. Что если у меня у одного элемента вес 5000, а у другого 1? Зачем мне ради одной операции строить массив длинной в 5001 элемент? :)

Конечно же, есть варианты гораздо проще и изящнее. Списывайте:


import random 

def weighted_choice(choices):
    u'''
    Взвешенный псевдослучайный выбор. Чем выше вес, тем выше шанс выпадения значения.
    @param choices: список или кортеж пар вида: (   ('choice', 14), ('choice2', 11), ... (<значение>, <вес>  ) 
    '''
    total = sum(w for c,w in choices)
    r = random.uniform(0, total)
    upto = 0
    for c, w in choices:
        if upto+w > r:
            return c
        upto += w
    assert False, "Shouldn't get here"

Работает вот так:

test_choices = ( ('Foo', 8), ('Bar', 3), ('god damn!', 1) )
for i in xrange(10):
    print weighted_choice(test_choices)

2 комментария:

  1. Вполне логичный способ, при больших значениях веса думаю разница в производительности будет очень заметна.

    ОтветитьУдалить
  2. К сожалению, многим людям он в голову не приходит :)

    ОтветитьУдалить