Premium

Кучи

Вес последнего камня

Каждый раз брать два самых тяжелых камня с помощью кучи

Легкий

Постановка задачи

Описание (LeetCode - 1046. Last Stone Weight)

Дан массив stones, где stones[i] - вес i-го камня.

На каждом шаге мы выбираем два самых тяжелых камня x и y, где x <= y, и сталкиваем их:

  • если x == y, оба камня уничтожаются;
  • если x < y, камень весом x уничтожается, а камень весом y получает вес y - x.

Нужно вернуть вес последнего оставшегося камня. Если камней не осталось, вернуть 0.

Пример 1:

stones = [2, 7, 4, 1, 8, 1]

Результат:

1

Пояснение:

Сначала сталкиваем камни 7 и 8, получаем 1. Массив становится [2, 4, 1, 1, 1].

Затем сталкиваем 4 и 2, получаем 2. Массив становится [2, 1, 1, 1].

Дальше сталкиваем 2 и 1, получаем 1. Массив становится [1, 1, 1].

Потом сталкиваем 1 и 1, оба камня уничтожаются. Остается один камень весом 1 - это и есть ответ.

Пример 2:

stones = [1]

Результат:

1

Ограничения:

  • 1 <= stones.length <= 30
  • 1 <= stones[i] <= 1000

Решение

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

Идея решения:

Положим все камни в максимальную кучу. Пока в куче есть хотя бы два камня, достаем два самых тяжелых. Если их веса отличаются, кладем обратно разницу.

В конце возможны два случая:

  1. В куче остался один камень - возвращаем его вес.
  2. Куча пуста - возвращаем 0.

Основные шаги:

  1. Создаем максимальную кучу.
  2. Добавляем в нее все камни.
  3. Пока в куче больше одного камня:
    • достаем самый тяжелый камень y;
    • достаем второй самый тяжелый камень x;
    • если y != x, добавляем y - x обратно.
  4. Возвращаем оставшийся камень или 0.

Рассмотрим пример:

stones = [2, 7, 4, 1, 8, 1]
  1. Берем 8 и 7, остается 1.
  2. Берем 4 и 2, остается 2.
  3. Берем 2 и 1, остается 1.
  4. Берем 1 и 1, оба уничтожаются.
  5. Остался камень весом 1.

Ответ: 1.

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

Оценка сложности

Временная сложность

Построение кучи занимает O(n). На каждом шаге мы делаем извлечение и, возможно, добавление за O(log n). Всего шагов не больше n.

Итоговая временная сложность: O(n log n).

Пространственная сложность

Куча хранит до n камней, поэтому дополнительная память равна O(n).

Код решения

Приведем код решения.

import (
    "github.com/emirpasic/gods/trees/binaryheap"
    "github.com/emirpasic/gods/utils"
)
 
func lastStoneWeight(stones []int) int {
    inverseIntComparator := func(a, b interface{}) int {
        return -utils.IntComparator(a, b)
    }
    heap := binaryheap.NewWith(inverseIntComparator)
 
    for _, stone := range stones {
        heap.Push(stone)
    }
 
    for heap.Size() > 1 {
        yVal, _ := heap.Pop()
        xVal, _ := heap.Pop()
        y := yVal.(int)
        x := xVal.(int)
        if y != x {
            heap.Push(y - x)
        }
    }
 
    if heap.Empty() {
        return 0
    }
    val, _ := heap.Pop()
    return val.(int)
}

Итоги

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

Войдите чтобы отмечать прогресс