Skip to content

Heap Basics

Table of Contents

1046. Last Stone Weight

"""
- Heap
    - Time: O(n log n); Space: O(n)
- 0/1 Knapsack
    - Time: O(n); Space: O(n)
"""

from heapq import heapify, heappop, heappush
from typing import List


# Heap
def lastStoneWeightHeap(stones: List[int]) -> int:
    maxHeap = [-s for s in stones]
    heapify(maxHeap)

    while len(maxHeap) > 1:
        s1 = heappop(maxHeap)
        s2 = heappop(maxHeap)

        if s1 != s2:
            heappush(maxHeap, s1 - s2)

    return -maxHeap[0] if maxHeap else 0


# 0/1 Knapsack
def lastStoneWeightKnapsack(stones: List[int]) -> int:
    total = sum(stones)
    target = total // 2

    dp = [0 for _ in range(target + 1)]

    for i in stones:
        for j in range(target, i - 1, -1):
            dp[j] = max(dp[j], dp[j - i] + i)

    return total - 2 * dp[target]


if __name__ == "__main__":
    stones = [2, 7, 4, 1, 8, 1]
    assert lastStoneWeightHeap(stones) == 1
    assert lastStoneWeightKnapsack(stones) == 1
#include <iostream>
#include <queue>
#include <vector>

using namespace std;

int lastStoneWeight(vector<int> &stones) {
    priority_queue<int> maxHeap(stones.begin(), stones.end());

    // Only while there are at least two stones to smash
    while (maxHeap.size() > 1) {
        int first = maxHeap.top();
        maxHeap.pop();
        int second = maxHeap.top();
        maxHeap.pop();

        if (first != second) {
            maxHeap.push(first - second);
        }
    }

    return maxHeap.empty() ? 0 : maxHeap.top();
}

int main() {
    vector<int> stones = {2, 7, 4, 1, 8, 1};
    cout << lastStoneWeight(stones) << endl;  // 1
    return 0;
}

3264. Final Array State After K Multiplication Operations I

import heapq
from typing import List


# Brute Force
def getFinalStateBF(nums: List[int], k: int, multiplier: int) -> List[int]:
    for _ in range(k):
        minNum = min(nums)
        idx = nums.index(minNum)
        nums[idx] *= multiplier

    return nums


# Heap
def getFinalStateHeap(nums: List[int], k: int, multiplier: int) -> List[int]:
    minHeap = []
    for idx, num in enumerate(nums):
        heapq.heappush(minHeap, (num, idx))

    for _ in range(k):
        num, idx = heapq.heappop(minHeap)
        nums[idx] = num * multiplier
        heapq.heappush(minHeap, (nums[idx], idx))

    return nums


k = 5
multiplier = 2
print(getFinalStateBF([2, 1, 3, 5, 6], k, multiplier))  # [8, 4, 6, 5, 6]
print(getFinalStateHeap([2, 1, 3, 5, 6], k, multiplier))  # [8, 4, 6, 5, 6]

2558. Take Gifts From the Richest Pile

from heapq import heapify, heappop, heappush
from math import isqrt
from typing import List


# Heap
def pickGifts(gifts: List[int], k: int) -> int:
    maxHeap = [-g for g in gifts]
    heapify(maxHeap)

    for _ in range(k):
        cur = heappop(maxHeap)

        if cur == -1:
            heappush(maxHeap, cur)
            break

        heappush(maxHeap, -isqrt(-cur))

    return sum(-i for i in maxHeap)


if __name__ == "__main__":
    assert pickGifts([25, 64, 9, 4, 100], 4) == 29
    assert pickGifts([1, 1, 1, 1], 4) == 0

2336. Smallest Number in Infinite Set

from heapq import heappop, heappush


# Heap
class SmallestInfiniteSet:
    def __init__(self):
        self.cur_min = 1
        self.added = set()
        self.min_heap = []

    def popSmallest(self) -> int:
        if self.min_heap:
            res = heappop(self.min_heap)
            self.added.remove(res)
            return res

        res = self.cur_min
        self.cur_min += 1
        return res

    def addBack(self, num: int) -> None:
        if num < self.cur_min and num not in self.added:
            self.added.add(num)
            heappush(self.min_heap, num)


if __name__ == "__main__":
    sis = SmallestInfiniteSet()
    assert sis.popSmallest() == 1
    sis.addBack(2)
    assert sis.popSmallest() == 2
    assert sis.popSmallest() == 3
    sis.addBack(1)
    assert sis.popSmallest() == 1
    assert sis.popSmallest() == 4
    sis.addBack(3)
    assert sis.popSmallest() == 3

2530. Maximal Score After Applying K Operations

from heapq import heapify, heappop, heappush
from math import ceil
from typing import List


# Heap
def maxKelements(nums: List[int], k: int) -> int:
    res = 0
    maxHeap = [-n for n in nums]
    heapify(maxHeap)

    while k > 0:
        cur = -heappop(maxHeap)
        res += cur
        heappush(maxHeap, -ceil(cur / 3))
        k -= 1

    return res


if __name__ == "__main__":
    assert maxKelements([10, 10, 10, 10, 10], 5) == 50
    assert maxKelements([1, 10, 3, 3, 3], 3) == 17
    assert maxKelements([1, 2, 3, 4, 5], 5) == 16

3066. Minimum Operations to Exceed Threshold Value II

from heapq import heapify, heappop, heappush
from typing import List


# Heap
def minOperations(nums: List[int], k: int) -> int:
    heapify(nums)
    res = 0

    while nums[0] < k:
        x = heappop(nums)
        y = heappop(nums)
        heappush(nums, x * 2 + y)
        res += 1

    return res


if __name__ == "__main__":
    assert minOperations([2, 11, 10, 1, 3], 10) == 2
    assert minOperations([1, 1, 2, 4, 9], 20) == 4

1962. Remove Stones to Minimize the Total

from heapq import heapify, heapreplace
from typing import List


# Heap
def minStoneSum(piles: List[int], k: int) -> int:
    maxHeap = [-p for p in piles]
    heapify(maxHeap)

    for _ in range(k):
        heapreplace(maxHeap, maxHeap[0] // 2)

    return -sum(maxHeap)


if __name__ == "__main__":
    assert minStoneSum([5, 4, 9], 2) == 12
    assert minStoneSum([4, 3, 6, 7], 3) == 12

703. Kth Largest Element in a Stream

  • LeetCode | 力扣

  • Tags: Tree, Design, Binary Search Tree, Heap Priority Queue, Binary Tree, Data Stream

import heapq
from typing import List


# Heap
class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.heap = []
        for num in nums:
            self.add(num)

    def add(self, val: int) -> int:
        heapq.heappush(self.heap, val)

        if len(self.heap) > self.k:
            heapq.heappop(self.heap)

        return self.heap[0]


obj = KthLargest(3, [4, 5, 8, 2])
print(obj.add(3))  # 4
print(obj.add(5))  # 5
print(obj.add(10))  # 5

3275. K-th Nearest Obstacle Queries

from heapq import heappop, heappush
from typing import List


# Heap
def resultsArray(queries: List[List[int]], k: int) -> List[int]:
    n = len(queries)
    res = [-1 for _ in range(n)]
    maxHeap = []

    for i in range(n):
        dist = abs(queries[i][0]) + abs(queries[i][1])
        heappush(maxHeap, -dist)

        if i < k - 1:
            continue

        while len(maxHeap) > k:
            heappop(maxHeap)

        res[i] = -maxHeap[0]

    return res


if __name__ == "__main__":
    queries = [[1, 2], [3, 4], [2, 3], [-3, 0]]
    k = 2
    assert resultsArray(queries, k) == [-1, 7, 5, 3]

2208. Minimum Operations to Halve Array Sum

import heapq
from typing import List


def halveArray(nums: List[int]) -> int:
    if not nums:
        return 0

    res = 0
    cur_sum = sum(nums)
    target = cur_sum / 2

    max_heap = [-num for num in nums]  # max heap
    heapq.heapify(max_heap)

    while cur_sum > target:
        mx = -heapq.heappop(max_heap)
        new = mx / 2
        heapq.heappush(max_heap, -new)
        cur_sum -= new
        res += 1

    return res


if __name__ == "__main__":
    assert halveArray([5, 19, 8, 1]) == 3
    assert halveArray([3, 8, 20]) == 3

2233. Maximum Product After K Increments

3296. Minimum Number of Seconds to Make Mountain Height Zero

from bisect import bisect_left
from heapq import heapify, heapreplace
from math import isqrt
from typing import List


# Min Heap
def minNumberOfSecondsMinHeap(
    mountainHeight: int, workerTimes: List[int]
) -> int:
    minHeap = [(t, t, t) for t in workerTimes]
    heapify(minHeap)

    for _ in range(mountainHeight):
        nxt, delta, base = minHeap[0]
        heapreplace(
            minHeap,
            (
                nxt + delta + base,
                delta + base,
                base,
            ),
        )
    return nxt


# Binary Search Min Answer
def minNumberOfSecondsBinarySearchMin(
    mountainHeight: int, workerTimes: List[int]
) -> int:
    def check(m: int) -> bool:
        left_h = mountainHeight
        for t in workerTimes:
            left_h -= (isqrt(m // t * 8 + 1) - 1) // 2
            if left_h <= 0:
                return True
        return False

    max_t = max(workerTimes)
    h = (mountainHeight - 1) // len(workerTimes) + 1
    return bisect_left(range(max_t * h * (h + 1) // 2), True, 1, key=check)


if __name__ == "__main__":
    mountainHeight = 4
    workerTimes = [2, 1, 1]
    assert minNumberOfSecondsMinHeap(mountainHeight, workerTimes) == 3
    assert minNumberOfSecondsBinarySearchMin(mountainHeight, workerTimes) == 3

1942. The Number of the Smallest Unoccupied Chair

1801. Number of Orders in the Backlog

2406. Divide Intervals Into Minimum Number of Groups

2462. Total Cost to Hire K Workers

1834. Single-Threaded CPU

3408. Design Task Manager

1792. Maximum Average Pass Ratio

2931. Maximum Spending After Buying Items

1882. Process Tasks Using Servers

2402. Meeting Rooms III

253. Meeting Rooms II

"""
- Given an array of meeting time `intervals` where `intervals[i] = [start_i, end_i]`, return the minimum number of conference rooms required.
"""

import heapq
from typing import List


# Heap
def minMeetingRooms(intervals: List[List[int]]) -> int:
    if not intervals:
        return 0

    intervals.sort(key=lambda x: x[0])
    minHeap = [intervals[0][1]]

    for i in range(1, len(intervals)):
        if intervals[i][0] >= minHeap[0]:
            heapq.heappop(minHeap)
        heapq.heappush(minHeap, intervals[i][1])

    return len(minHeap)


if __name__ == "__main__":
    intervals = [[0, 30], [5, 10], [15, 20]]
    assert minMeetingRooms(intervals) == 2

1167. Minimum Cost to Connect Sticks

from heapq import heapify, heappop, heappush
from typing import List


# Heap
def connectSticks(sticks: List[int]) -> int:
    n = len(sticks)
    heapify(sticks)
    res = 0

    while n > 1:
        x = heappop(sticks)
        y = heappop(sticks)
        res += x + y
        heappush(sticks, x + y)
        n -= 1

    return res


if __name__ == "__main__":
    assert connectSticks([2, 4, 3]) == 14
    assert connectSticks([1, 8, 3, 5]) == 30
    assert connectSticks([5]) == 0
    assert connectSticks([1, 2, 3, 4, 5]) == 33
    assert connectSticks([1, 1, 1]) == 5

Comments