Разбор популярной задачи уровня Leetcode Medium из бигтехов

Я знаю 4 компании из tier 1-2, где её давали на собесах в разных вариациях, поэтому пост обещает быть полезным. Погнали🚀

Более того, я хочу не просто показать вам код, а описать путь к оптимальному решению, каким он может быть на реальном собеседовании (см. пост про алгоритмическую сложность).

Пусть дан массив пар с городами и их населением, например:
[['Kaliningrad', 488690], ['Sochi', 445149], ['Kolomna', 213701]]

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

Интуиция: давайте положим в корзину 488690 бумажек с надписью «Калининград», 445149 — с надписью «Сочи», 213701 — с надписью «Коломна», и вытянем случайную из них.
Это ровно то, что требуется сделать по условию, и чисто технически мы действительно могли бы создать такой огромный список строк (но не set) с названиями городов, и выбирать случайную.

А надо ли тратить столько памяти (пожалейте деревья)?
Окей, давайте заметим, что можно сэмплить случайное число между 1 и суммой значений в словаре. Если выпало от 1 до 488690, то возвращаем Калининград, от 488691 до 488690 + 445149 — Сочи, от 488691 + 445150 до 488690 + 445149 + 213701 — Коломну. Мы избавились от вспомогательного массива, который мог бы оказаться большим.

Тем не менее, а если городов очень много, удобно ли работать с такими большими целыми? Давайте нормируем так, чтобы сумма всех жителей упомянутых городов равнялась 1.0 — это лучше соответствует определению вероятности в принципе, а работа с float достаточно точная на практике. Делим всё на 488690 + 445149 + 213701 и вызываем функцию для случайного числа на полуинтервале [0, 1) — в питоне это random.random().

Опишем алгоритмически: за один проход по массиву (O(N), где N — количество городов) мы формируем массив кумулятивных сумм вероятностей выпадения каждого города:

cur = 0.0

sums = []

total = sum([city[1] for city in cities])

for city in cities:

    sums.append(cur)

    cur += city[1] / total

assert cur == 1.0 # должно сойтись по построению



# print(sums)

# [0.0, 0.4258587935932516, 0.8137746832354428]

Далее вызываем рандомное действительное число и ищем слева направо первый элемент массива sums, который больше него. Тогда ответ — предыдущий город🌇

r = random.random()

for i in range(len(sums)):

    if r < sums[i]:

        return cities[i - 1][0]

return cities[-1][0]

Выхода за границу массива слева не будет из-за строгого <, а последний return отвечает за случай r == 1.0 (вообще такого не бывает, но чтобы точно собеседующий был доволен).

Это ещё один проход по массиву размера N (aka линейный поиск), то есть сложность алгоритма O(N) + O(N) = O(N).

Высший пилотаж: доводим до оптимального решения

Выше я не просто так упомянула сложность каждого из шагов.

Вторая часть может быть оптимизирована до O(log N), если линейный поиск заменить на бинарный. Это можно сделать, так как выполнено основное требование — массив кумулятивных сумм вероятностей является отсортированным по построению.

Теперь заметим, что по сути первое действие достаточно сделать всего один раз, а не каждый раз при вызове функции выбора случайного города (вызывать её будут много раз, раз речь зашла о распределении). Тогда можно вместо функции написать класс, где эта часть будет вынесена в __init__, а основной метод sample() будет быстрее.

Это в целом хорошая практика — выносить предвычисления из основного метода, так как за много вызовов sample() один целый проход по массиву амортизируется («размазывается» по всем вызовам).

Итоговый код приложу в комментарии! Как вам такой формат?

И голосуйте: — если переслал другу, который ищет работу, — слишком сложно, до свидания, — наконец-то понятно, спасибо

#найм_и_собесы #хардов_пост
@analytess