Простой алгоритм случайной выборки с учетом веса

Иногда вам может понадобиться выбрать случайный элемент из списка с учетом того, что некоторые элементы имеют больший шанс выбора, чем другие (имеют больший «вес»). Например, вы можете взять список приложений и количество загрузок, и случайным образом выбрать «Популярное приложение» в зависимости от количества загрузок.

В этой заметке я покажу вам два подхода к «взвешенному» случайному выбору — один подходит для небольших списков, а другой оптимизирован для большего числа элементов.

Простой алгоритм случайной выборки с учетом веса

В общем виде этот алгоритм можно описать так:

1. Выбрать случайное число между единицей и суммой «весов» всех элементов
2. Спускаться по списку элементов добавляя к счетчику вес текущего элемента
3. Проверить, если счетчик (шаг №2) больше или равен случайному числу (шаг №1), то закончить цикл и вернуть текущий элемент. В противном случае перейдите к шагу №2.

Этот алгоритм прост в реализации и довольно быстр, когда число элементов не велико, или когда вам нужно сделать выбор один раз. Ниже приводится функция, которая принимает массив элементов для выбора, а также массив соответствующих им весов, и возвращает случайно выбранный элемент из первого массива. Вы можете использовать любое целое положительное число, как вес.

/**
 * Выборка случайного элемента с учетом веса
 *
 * @param array $values индексный массив элементов
 * @param array $weights индексный массив соответствующих весов
 * @return mixed выбранный элемент
 */
function weighted_random_simple ( $values, $weights )
{
    $total = array_sum( $weights );
    $n = 0;
    $num = mt_rand( 1, $total );
    foreach ( $values as $i => $value )
    {
        $n += $weights[$i];
        if ( $n >= $num )
        {
            return $values[$i];
        }
    }
}

Вот пример скрипта который выведет либо A, B, C или с вероятностью 15%, 35% и 50% соответственно:

$values = array('A', 'B', 'C');
$weights = array(3, 7, 10);
echo weighted_random_simple($values, $weights);

Алгоритм случайной выборки из тысяч элементов

Описанный выше алгоритм может работать очень медленно, когда список элементов велик, и вам необходимо сделать несколько выборок. Это потому, что он должен пройтись по всему массиву каждый раз при обращении к функции.

Однако, алгоритм может быть расширен, чтобы сделать его значительно быстрее. Вместо вычисления общего веса (шаг №1) и счетчика (шаг №2) каждый раз, можно сделать это один раз и сохранить значения счетчиков в массиве. Тогда мы сможем использовать бинарный поиск, чтобы быстро выбрать правильный элемент. Ниже приведен модифицированный вариант функции:

/**
 * Случайно выбирает один из элементов на основе их веса.
 * Оптимизирован для большого числа элементов.
 *
 * @param array $values индексный массив элементов
 * @param array $weights индексный массив соответствующих весов
 * @param array $lookup отсортированный массив для поиска
 * @param int $total_weight сумма всех весов
 * @return mixed выбранный элемент
 */
function weighted_random($values, $weights, $lookup = null, $total_weight = null)
{
    if ($lookup == null OR $total_weight == null)
    {
        list($lookup, $total_weight) = calc_lookups($values, $weights);
    }
    $r = mt_rand(1, $total_weight);
    return $values[binary_search($r, $lookup)];
}
/**
 * Создание массива используемого в бинарном поиске
 *
 * @param array $values
 * @param array $weights
 * @return array
 */
function calc_lookups($values, $weights)
{
    $lookup = array();
    $total_weight = 0;
    for ($i=0; $i < count($weights); $i++)
    {
        $total_weight += $weights[$i];
        $lookup[$i] = $total_weight;
    }
    return array($lookup, $total_weight);
}
/**
 * Ищет в массиве элемент по номеру и возвращает элемент если он найден.
 * В противном случае возвращает позицию, где он должен быть вставлен,
 * или count($haystack)-1, если $needle больше чем любой элемент в массиве.
 *
 * @param int $needle
 * @param array $haystack
 * @return int
 */
function binary_search($needle, $haystack)
{
    $high = count($haystack) - 1;
    $low = 0;
    while ( $low < $high )
    {
        $probe = (int)(($high + $low) / 2);
        if ($haystack[$probe] < $needle)
        {
            $low = $probe + 1;
        }
        elseif ($haystack[$probe] > $needle)
        {
            $high = $probe - 1;
        }
        else
        {
            return $probe;
        }
    }
    if ( $low != $high )
    {
        return $probe;
    }
    else
    {
        return ($haystack[$low] >= $needle) ? $low : $low + 1;
    }
}

Описанный выше скрипт также содержит две новые функции — calc_lookups которая вычисляет массив для использования в бинарном поиске, и непосредственно функция binary_search которая осуществляет бинарный поиск. Пример использования скрипта:

// Рассчет массивов (1 раз)
list($lookup, $total_weight) = calc_lookups($values, $weights);
//....
// Каждый раз когда вам необходимо выбрать случайный элемент:
$val = weighted_random($values, $weights, $lookup, $total_weight);

В заключение

Чтобы дать вам представление о том, какова скорость этих алгоритмов: Для каждого из них я использовал массив включающий 10 000 элементов, 10 000 раз подряд. Первый алгоритм отработал за 13 секунд, а второй всего 0,09 секунд.