Создание взвешенной версии random.choice в Python

Быстрый ответ

Если перед вами стоит задача выбора элемента на основе его веса, вы можете воспользоваться функцией random.choices() с параметром weights, управляющим вероятностным распределением. Код ниже демонстрирует простой способ выбрать случайный элемент, учитывая его вес:

import random

элементы = ['A', 'B', 'C', 'D']
веса = [10, 1, 1, 1]

выбранный_элемент = random.choices(элементы, веса)[0]
print(выбранный_элемент)

Заметьте, что 'A' имеет наибольший вес, а random.choices() возвращает список, так что мы используем индекс [0], чтобы получить конкретный элемент.

Для работы с большим объемом данных: используйте NumPy

Если вы работаете с большим объемом данных или вам требуется максимальная точность, NumPy с функцией numpy.random.choice() сможет вам помочь. Задайте вероятности через параметр p и укажите, планируете ли вы использовать выборку с заменой с помощью атрибута replace.

Пример кода:

import numpy as np

элементы = ['A', 'B', 'C', 'D']
веса = np.array([10, 1, 1, 1])
веса = веса / np.sum(веса)  # Нормализация весов

выбранный_элемент = np.random.choice(элементы, p=веса)
print(выбранный_элемент)

Непривычное использование: bisect для работы с накопленными весами

Мало кто знает, что модуль bisect в Python тоже может быть эффективно использован для взвешенного выбора. Он особенно полезен при работе с накопленными весами:

import random
import bisect

элементы = ['A', 'B', 'C', 'D']
веса = [10, 20, 30, 40]

# Кумулятивное распределение веса
накопленные_веса = [sum(веса[:i+1]) for i in range(len(веса))]

# Случайное число в пределах суммы весов
случайное_число = random.uniform(0, накопленные_веса[-1])

# Определение выбранного элемента
выбранный_элемент = элементы[bisect.bisect(накопленные_веса, случайное_число)]
print(выбранный_элемент)

Больше читабельности с помощью zip: скрытая возможность Python

Можно упростить понимание кода и сделать его более читабельным, используя функцию zip() для сопоставления элементов с их весами:

def взвешенный_случайный_выбор(пары):
    общий_вес = sum(вес for элемент, вес in пары)
    r = random.uniform(0, общий_вес)
    добавка = 0
    for элемент, вес in пары:
        добавка += вес
        if добавка >= r:
            return элемент

из_пар = zip(элементы, веса)
print(взвешенный_случайный_выбор(из_пар))

Визуализация

Давайте рассмотрим пример с шариками разного цвета, каждый из которых имеет свой вес:

Шарик   | Вес (Вероятность)
----------------------------------------
🎈Красный | 10 🏋️‍♂️
🎈Синий   | 5 🏋️‍♂️
🎈Зелёный | 2 🏋️‍♂️
🎈Жёлтый  | 1 🏋️‍♂️

Используя random.choices(), можно сказать, что выбор определенного шарика напоминает выбор шарика сильным ветром, при котором большей вероятности подвержены шарики с большим весом:

import random
выбор = random.choices(
  ["🎈Красный", "🎈Синий", "🎈Зелёный", "🎈Жёлтый"],
  [10, 5, 2, 1]
)

Так, тяжёлые шарики вероятнее всех останутся на месте:

Случайный ветерок 🌬️ -> [🎈Красный, 🎈Красный, 🎈Синий, 🎈Красный]
# Красный шарик — самый "тяжелый"

Работаем без циклов: используем NumPy

Если вам необходим всего один случайный элемент из набора, использование циклов может оказаться избыточным. В таких случаях оптимальнее воспользоваться функцией, способной выполнить задачу одной командой, вроде np.random.choice из NumPy для выборки без повторений.

Применяем все типы весов

Наши функции должны корректно работать с любыми числовыми типами весов: целыми числами, дробями или даже комплексными числами, если вы пожелаете их использовать.

Погружаемся в документацию Python

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

Уникальный выбор с помощью replace в NumPy

Иногда важно, чтобы вы выбрали каждый элемент только один раз. В этом случае установите параметр replace функции numpy.random.choice равным False, что обеспечит уникальность выбора:

import numpy as np

элементы = np.array(['A', 'B', 'C', 'D'])
веса = np.array([10, 20, 30, 40])

# При replace=False каждый элемент можно выбрать только один раз
уникальные_выборы = np.random.choice(элементы, size=2, replace=False, p=веса/веса.sum())
print(уникальные_выборы)

Полезные материалы

1. Документация Python 3.12.1 – random.choices() — генерация взвешенных случайных выборок.
2. Документация NumPy v1.26 – numpy.random.choice() — для сложных случаев случайного выбора с учетом веса.
3. Взвешенная версия random.choice на Stack Overflow — советы и решения от сообщества разработчиков.
4. Medium – Взвешенный случайный выбор в Python — инструменты для реализации взвешенного выбора с учетом особенностей работы с большими данными.
5. Блог Элай Бендерски – Генерация взвешенных случайных чисел в Python — различные алгоритмы выбора.
6. Educative: Интерактивные курсы – Модуль Random в Python — руководство по использованию модуля random.