Боря программирует

Weighted sampling without replacement



Вернёмся к скучным постам про алгоритмы. Допустим у нас есть N объектов, где объекту i сопоставлен какой-то вес w[i]. Мы хотим сгенерировать перестановку из этих элементов в соответствии с весами. Изначально возьмем пустой массив и будем добавлять в него элементы. Каждый раз добавляем случайный элемент, который еще не добавлен в массив. При этом вероятность взять элемент i должна быть пропорциональна w[i]. Как это сделать быстрее чем за O(N^2)?



Как спортивный программист я умею это делать с помощью дерева отрезков на сумму. Изначально запишем в ячейку i число w[i]. Чтобы выбрать очередной объект, посчитаем текущую сумму в дереве отрезков S, сгенерируем случайное число R от 1 до S. А потом c помощью спуска по дереву найдём наименьший префикс p такой, что сумма в ячейках 1..p хотя бы R. Добавим число p в ответ, а в дереве отрезков запишем туда 0 вместо w[p]. Повторим так N раз. Суммарно это работает за O(N log N).



Этот способ работает, но он довольно сложный. Оказывается можно сделать все тоже самое гораздо быстрее и за несколько строк кода.



Рассмотрим следующий алгоритм. Для каждого объекта i сгенируем случайное вещественные число r[i] от 0 до w[i]. А потом отсортируем все объекты в порядке убывания r[i]. Утверждается, что это и есть перестановка, которую мы хотели найти.



Например, если объекты i и j имели одинаковые веса, то i окажется раньше j в финальной перестановке с вероятностью 50%, как и нужно.



Если объект i имел вес больше чем j, то в среднем r[i] окажется больше чем r[j], и i в среднем окажется раньше в перестановке чем j, как и нужно.



К сожалению, если быть более внимательным, и честно посчитать вероятности получения нужных перестановок в ответе, то оказывается, что алгоритм работает неправильно. Это можно понять даже на примере из двух объектов.



Но оказывается, что можно немного изменить способ генерирования r[i], и метод начинает правильно работать! Расскажите в комментариях как исправить решение.