1. Условие задачи: камнем по голове
В черном цилиндре, наряде старинном путник на праздник очень спешит. По горам пробирается и улыбается, но лавина из камней срывается в пропасть с горных вершин.
Пока лавина несётся вниз, два самых тяжёлых камня раз за разом сталкиваются друг с другом. Пусть их веса равны x и y, где x <= y.
- Если
x == y, камни разбиваются вдребезги, и не остаётся ни одного. - Если
x < y, меньший камень исчезает, а от большего остаётся обломок весомy - x.
Затем среди уцелевших камней снова сталкиваются два самых тяжёлых. Так продолжается, пока не останется не больше одного камня.
Дано: массив stones, где stones[i] обозначает вес очередного камня.
Нужно: вернуть вес последнего камня, который долетит до головы путника. Если все камни уничтожили друг друга, вернуть 0.
Пример
Для stones = [2, 7, 4, 1, 8, 1] столкновения могут выглядеть так:
8и7превращаются в камень весом1.4и2превращаются в камень весом2.2и1превращаются в камень весом1.1и1уничтожают друг друга.- Остаётся камень весом
1.
Ответ: 1.
Ограничения: от 1 до 30 камней, вес каждого составляет от 1 до 1000.
2. Решение через сортировку
Нам всё время нужны два самых тяжёлых камня. Самый прямой способ их найти: отсортировать массив. Тогда два нужных камня окажутся в конце списка, и мы сможем забрать их с помощью pop().
После столкновения может появиться новый камень весом y - x. Просто добавим его обратно в список. Но порядок после этого снова нарушится, поэтому на следующем круге опять выполним сортировку.
Алгоритм по шагам
- Пока в списке есть хотя бы два камня, сортируем его по возрастанию.
- Достаём с конца самый тяжёлый камень
y. - Достаём следующий по весу камень
x. - Если их веса различаются, добавляем обратно обломок
y - x. - Если список опустел, возвращаем
0. Иначе возвращаем вес последнего камня.
Код на Python
class Solution:
def lastStoneWeight(self, stones: List[int]) -> int:
while len(stones) > 1:
stones.sort()
y = stones.pop() # Самый тяжёлый камень
x = stones.pop() # Второй по тяжести
if x != y:
stones.append(y - x)
return stones[0] if stones else 0
Сложность сортировки
За одно столкновение число камней уменьшается хотя бы на один, поэтому столкновений будет не больше N - 1.
Перед каждым из них мы сортируем список. Одна сортировка стоит O(N log N), а повторяем мы её O(N) раз. Поэтому верхняя оценка времени составляет O(N² log N).
3. Оптимальное решение: куча
Сортировка делает лишнюю работу: на каждом шаге она полностью упорядочивает все камни, хотя нам нужны только два самых тяжёлых.
Для таких случаев существует куча, или priority queue. Пока не будем вникать в её внутреннее устройство – можно относится к ней как к хэш-таблице (словарю) – мы не пишем его сами и можем не знать как он работает – просто бессовестно пользуемся теми преимуществами, которые он нам даёт и всё. Нам достаточно знать три свойства кучи:
- кучу можно построить из массива за
O(N); - самый приоритетный элемент можно достать за
O(log N); - новый элемент можно добавить за
O(log N).
В этой задаче приоритетнее камень с наибольшим весом, то есть нам нужна max-heap. Но стандартный модуль Python heapq умеет быстро доставать минимум. Превратим минимум в максимум простым трюком: будем хранить все веса со знаком минус.
Например, веса [2, 7, 4] превратятся в [-2, -7, -4]. Самое маленькое число здесь -7, но оно соответствует самому тяжёлому камню весом 7.
Алгоритм по шагам
- Меняем знак каждого веса и строим из полученного списка min-heap.
- Пока в куче есть хотя бы два камня, дважды достаём минимум.
- Возвращаем числам обычные знаки и получаем два самых тяжёлых камня
yиx. - Если
x != y, кладём в кучу обломок весомy - x, снова со знаком минус. - Если куча опустела, возвращаем
0. Иначе меняем знак последнего элемента и возвращаем его вес.
Визуализация на stones = [2, 7, 4, 1, 8, 1]
Для понимания в таблице показаны обычные положительные веса, хотя внутри кучи они хранятся со знаком минус.
| Шаг | Два самых тяжёлых | Результат столкновения | Оставшиеся веса |
|---|---|---|---|
| 1 | 8 и 7 |
обломок 1 |
[4, 2, 1, 1, 1] |
| 2 | 4 и 2 |
обломок 2 |
[2, 1, 1, 1] |
| 3 | 2 и 1 |
обломок 1 |
[1, 1, 1] |
| 4 | 1 и 1 |
оба уничтожены | [1] |
До головы путника долетит камень весом 1 – цилиндр чёрный будет смят гормошкой(.
Код на Python
import heapq
class Solution:
def lastStoneWeight(self, stones: List[int]) -> int:
# heapq достаёт минимум, поэтому меняем знаки весов
heap = [-stone for stone in stones]
heapq.heapify(heap)
while len(heap) > 1:
y = -heapq.heappop(heap) # Самый тяжёлый камень
x = -heapq.heappop(heap) # Второй по тяжести
if x != y:
heapq.heappush(heap, -(y - x))
return -heap[0] if heap else 0
Почему куча быстрее
Построение кучи занимает O(N). Затем происходит не больше N - 1 столкновений. На каждом шаге мы выполняем два извлечения и не больше одной вставки, каждая из которых стоит O(log N).
Итог по времени: O(N log N).
Итог по памяти: O(N), потому что мы создаём отдельный список отрицательных весов.
По сравнению с самой простой симуляцией, которая заново сортирует список после каждого столкновения, мы улучшаем время с O(N² log N) до O(N log N).
Если сравнивать с более аккуратным решением через один раз отсортированный список и bisect.insort, улучшение будет с O(N²) до O(N log N). Никакого O(N log K) здесь нет: размер кучи в худшем случае остаётся порядка N, поэтому каждая операция с ней стоит O(log N).
Краевые случаи
- Один камень: столкновений нет, сразу возвращаем его вес.
- Два одинаковых камня: они уничтожают друг друга, возвращаем
0. - Два разных камня: остаётся их разность.
- Все камни одинаковые: они попарно исчезают; ответ зависит от чётности их количества.
- После столкновения появился очень лёгкий обломок: куча сама поставит его на нужное место, искать позицию вручную не нужно.
Вывод
Ключевой паттерн задачи: если много раз нужно доставать самый большой или самый маленький элемент, стоит подумать о куче.
Обычная сортировка проста и хорошо объясняет симуляцию, но каждый раз упорядочивает весь список заново. Куча хранит только необходимый порядок приоритетов и потому сокращает время до O(N log N).
Так камни продолжают биться друг о друга, пока над дорогой не повисает тишина. И только один вопрос остаётся без ответа: будут ли странно смотреть на путника местные собаки?
🐆