Heap — Worked LC Examples

# python
# LC 1046 - Last Stone Weight
# IDEA: max-heap (negated) of every stone; smash the top two, push back the difference
# time = O(n log n), space = O(n)
import heapq

def lastStoneWeight(stones):
    heap = [-s for s in stones]          # negated: heaviest on top
    heapq.heapify(heap)
    while len(heap) > 1:
        first = -heapq.heappop(heap)     # heaviest
        second = -heapq.heappop(heap)    # second heaviest
        if first != second:
            heapq.heappush(heap, -(first - second))
    return -heap[0] if heap else 0

// java
// LC 1046 - Last Stone Weight
// IDEA: max-heap of every stone; smash the top two, offer back the difference
// time = O(n log n), space = O(n)
public int lastStoneWeight(int[] stones) {
    PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
    for (int stone : stones) {
        maxHeap.offer(stone);
    }
    while (maxHeap.size() > 1) {
        int first = maxHeap.poll();          // heaviest
        int second = maxHeap.poll();         // second heaviest
        if (first != second) {
            maxHeap.offer(first - second);   // the survivor goes back into the pool
        }
    }
    return maxHeap.isEmpty() ? 0 : maxHeap.peek();
}

# python
# LC 703 - Kth Largest Element in a Stream
# IDEA: a min-heap of the k largest seen so far, kept between calls; its top is the kth largest
# time = O(log k) per add (O(n log k) to build), space = O(k)
import heapq

class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.heap = []                   # the k largest seen so far
        for num in nums:
            self.add(num)

    def add(self, val):
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)     # evict the smallest; the k largest remain
        return self.heap[0]              # smallest of the k largest = kth largest

// java
// LC 703 - Kth Largest Element in a Stream
// IDEA: a min-heap of the k largest seen so far, kept between calls; its top is the kth largest
// time = O(log k) per add, space = O(k)
class KthLargest {
    private final PriorityQueue<Integer> minHeap = new PriorityQueue<>();   // the k largest
    private final int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        for (int num : nums) {
            add(num);
        }
    }

    public int add(int val) {
        minHeap.offer(val);
        if (minHeap.size() > k) {
            minHeap.poll();                  // evict the smallest; the k largest remain
        }
        return minHeap.peek();
    }
}

# python
# LC 1167 - Minimum Cost to Connect Sticks
# IDEA: Huffman merge - pop the two shortest, pay their sum, push the merged stick back
# time = O(n log n), space = O(n)
import heapq

def connectSticks(sticks):
    heap = list(sticks)
    heapq.heapify(heap)
    cost = 0
    while len(heap) > 1:
        a = heapq.heappop(heap)          # the two shortest sticks
        b = heapq.heappop(heap)
        cost += a + b
        heapq.heappush(heap, a + b)      # the merged stick goes back into the pool
    return cost

# python
# LC 1492 - The kth Factor of n
# IDEA: enumerate divisor pairs up to sqrt(n); a size-k max-heap keeps the k smallest
# time = O(sqrt(n) log k), space = O(k)
import heapq

def kthFactor(n, k):
    heap = []                            # max-heap (negated) of the k smallest divisors
    def push(d):
        heapq.heappush(heap, -d)
        if len(heap) > k:
            heapq.heappop(heap)          # drop the largest; the k smallest remain
    i = 1
    while i * i <= n:
        if n % i == 0:
            push(i)
            if i != n // i:
                push(n // i)             # the paired divisor
        i += 1
    return -heap[0] if len(heap) == k else -1

# python
# LC 973 - K Closest Points to Origin
# IDEA: size-k max-heap (negated) by squared distance; evict the farthest when it passes k
# time = O(n log k), space = O(k)
import heapq

def kClosest(points, k):
    heap = []                            # max-heap (negated distance) of the k closest
    for x, y in points:
        dist = x * x + y * y             # squared distance: no sqrt, no float
        heapq.heappush(heap, (-dist, x, y))
        if len(heap) > k:
            heapq.heappop(heap)          # evict the farthest of the k + 1
    return [[x, y] for _, x, y in heap]

// java
// LC 973 - K Closest Points to Origin
// IDEA: size-k max-heap by squared distance; evict the farthest when it passes k
// time = O(n log k), space = O(k)
public int[][] kClosest(int[][] points, int k) {
    // max-heap by squared distance: the farthest of the k kept is on top
    PriorityQueue<int[]> maxHeap = new PriorityQueue<>(
        (a, b) -> Integer.compare(b[0] * b[0] + b[1] * b[1], a[0] * a[0] + a[1] * a[1]));
    for (int[] point : points) {
        maxHeap.offer(point);
        if (maxHeap.size() > k) {
            maxHeap.poll();                  // evict the farthest
        }
    }
    int[][] result = new int[k][];
    for (int i = 0; i < k; i++) {
        result[i] = maxHeap.poll();
    }
    return result;
}

# python
# LC 451 - Sort Characters By Frequency
# IDEA: count, then a max-heap by count; pop each character and emit its whole run
# time = O(n + d log d), space = O(n)   d = distinct characters
import heapq
from collections import Counter

def frequencySort(s):
    count = Counter(s)
    heap = [(-freq, ch) for ch, freq in count.items()]   # most frequent on top
    heapq.heapify(heap)
    parts = []
    while heap:
        neg_freq, ch = heapq.heappop(heap)
        parts.append(ch * -neg_freq)     # emit the whole run at once
    return "".join(parts)

// java
// LC 451 - Sort Characters By Frequency
// IDEA: count, then a max-heap of characters by count; pop each and emit its whole run
// time = O(n + d log d), space = O(n)   d = distinct characters
public String frequencySort(String s) {
    Map<Character, Integer> freq = new HashMap<>();
    for (char c : s.toCharArray()) {
        freq.put(c, freq.getOrDefault(c, 0) + 1);
    }
    PriorityQueue<Character> maxHeap = new PriorityQueue<>(
        (a, b) -> Integer.compare(freq.get(b), freq.get(a)));   // most frequent on top
    maxHeap.addAll(freq.keySet());
    StringBuilder sb = new StringBuilder();
    while (!maxHeap.isEmpty()) {
        char c = maxHeap.poll();
        for (int i = 0; i < freq.get(c); i++) {
            sb.append(c);                    // emit the whole run at once
        }
    }
    return sb.toString();
}

# python
# LC 1481 - Least Number of Unique Integers after K Removals
# IDEA: min-heap of counts; remove whole values, cheapest first, while k covers them
# time = O(n + d log d), space = O(d)   d = distinct values
import heapq
from collections import Counter

def findLeastNumOfUniqueInts(arr, k):
    heap = list(Counter(arr).values())   # min-heap of counts: cheapest to wipe out on top
    heapq.heapify(heap)
    while heap and heap[0] <= k:
        k -= heapq.heappop(heap)         # remove every copy of that value
    return len(heap)

# python
# LC 1647 - Minimum Deletions to Make Character Frequencies Unique
# IDEA: max-heap of counts; on a tie with the next count, delete one and push it back
# time = O(n + K log K), space = O(K)   K = distinct letters (26), each count shaved at most n times
import heapq
from collections import Counter

def minDeletions(s):
    heap = [-f for f in Counter(s).values()]   # max-heap of counts
    heapq.heapify(heap)
    deletions = 0
    while len(heap) > 1:
        top = -heapq.heappop(heap)
        if top == -heap[0]:              # the next count is the same: shave one off this one
            deletions += 1
            if top - 1 > 0:
                heapq.heappush(heap, -(top - 1))
        # otherwise `top` is now unique and leaves the heap for good
    return deletions

// java
// LC 1647 - Minimum Deletions to Make Character Frequencies Unique
// IDEA: max-heap of counts; on a tie with the next count, delete one and offer it back
// time = O(n + K log K), space = O(K)   K = distinct letters (26)
public int minDeletions(String s) {
    int[] freq = new int[26];
    for (char c : s.toCharArray()) {
        freq[c - 'a']++;
    }
    PriorityQueue<Integer> pq = new PriorityQueue<>(Collections.reverseOrder());   // max-heap of counts
    for (int f : freq) {
        if (f > 0) {
            pq.add(f);
        }
    }
    int deletions = 0;
    while (pq.size() > 1) {
        int top = pq.poll();
        if (top == pq.peek()) {              // tied with the next count: shave one off
            deletions++;
            if (top - 1 > 0) {
                pq.add(top - 1);
            }
        }                                    // otherwise top is unique and leaves for good
    }
    return deletions;
}

# python
# LC 373 - Find K Pairs with Smallest Sums
# IDEA: k-way merge over rows: seed (i, 0) for the first k rows, pop k times, advance the popped row
# time = O(k log k), space = O(k)
import heapq

def kSmallestPairs(nums1, nums2, k):
    # one source per row i: the pairs (i, 0), (i, 1), ... are sorted by sum
    heap = [(nums1[i] + nums2[0], i, 0) for i in range(min(k, len(nums1)))]
    heapq.heapify(heap)
    res = []
    while heap and len(res) < k:
        _, i, j = heapq.heappop(heap)
        res.append([nums1[i], nums2[j]])
        if j + 1 < len(nums2):
            heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))   # advance row i
    return res

// java
// LC 373 - Find K Pairs with Smallest Sums
// IDEA: k-way merge over rows: seed (i, 0) for the first k rows, poll k times, advance the polled row
// time = O(k log k), space = O(k)
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
    List<List<Integer>> res = new ArrayList<>();
    // {sum, i, j}: one entry per row i, the row's current head (i, j)
    PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> Integer.compare(a[0], b[0]));
    for (int i = 0; i < nums1.length && i < k; i++) {
        minHeap.offer(new int[]{nums1[i] + nums2[0], i, 0});   // (i, 0) is row i's smallest
    }
    while (res.size() < k && !minHeap.isEmpty()) {
        int[] cur = minHeap.poll();
        int i = cur[1], j = cur[2];
        res.add(Arrays.asList(nums1[i], nums2[j]));
        if (j + 1 < nums2.length) {
            minHeap.offer(new int[]{nums1[i] + nums2[j + 1], i, j + 1});   // advance row i
        }
    }
    return res;
}

# python
# LC 378 - Kth Smallest Element in a Sorted Matrix
# IDEA: k-way merge of the rows; pop k - 1 values, the top is then the kth smallest
# time = O(k log n), space = O(n)
import heapq

def kthSmallest(matrix, k):
    n = len(matrix)
    heap = [(matrix[r][0], r, 0) for r in range(min(k, n))]   # head of each row
    heapq.heapify(heap)
    for _ in range(k - 1):               # pop k - 1 values; the next top is the kth
        _, r, c = heapq.heappop(heap)
        if c + 1 < n:
            heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
    return heap[0][0]

// java
// LC 378 - Kth Smallest Element in a Sorted Matrix
// IDEA: k-way merge of the rows; poll k - 1 values, the top is then the kth smallest
// time = O(k log n), space = O(n)
public int kthSmallest(int[][] matrix, int k) {
    int n = matrix.length;
    // {value, row, col}: the current head of each row
    PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> Integer.compare(a[0], b[0]));
    for (int r = 0; r < n && r < k; r++) {
        minHeap.offer(new int[]{matrix[r][0], r, 0});
    }
    for (int i = 0; i < k - 1; i++) {        // poll k - 1 values; the next top is the kth
        int[] cur = minHeap.poll();
        int r = cur[1], c = cur[2];
        if (c + 1 < n) {
            minHeap.offer(new int[]{matrix[r][c + 1], r, c + 1});
        }
    }
    return minHeap.peek()[0];
}

# python
# LC 1675 - Minimize Deviation in Array
# IDEA: normalise odds up (x2); max-heap of current values; halve the max while it is even,
#       tracking the running min; the best max - min seen is the answer
# time = O(n log n log M), space = O(n)   M = the largest value
import heapq

def minimumDeviation(nums):
    heap = []                            # max-heap (negated) of each number's current value
    low = float('inf')
    for x in nums:
        if x % 2 == 1:
            x *= 2                       # an odd number's largest reachable value
        heapq.heappush(heap, -x)
        low = min(low, x)
    best = float('inf')
    while True:
        high = -heapq.heappop(heap)
        best = min(best, high - low)
        if high % 2 == 1:                # the max cannot shrink any further
            return best
        high //= 2
        low = min(low, high)
        heapq.heappush(heap, -high)

// java
// LC 1675 - Minimize Deviation in Array
// IDEA: normalise odds up (x2); max-heap of current values; halve the max while it is even
// time = O(n log n log M), space = O(n)   M = the largest value
public int minimumDeviation(int[] nums) {
    PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
    int min = Integer.MAX_VALUE;
    for (int x : nums) {
        if (x % 2 == 1) {
            x *= 2;                          // an odd number's largest reachable value
        }
        maxHeap.offer(x);
        min = Math.min(min, x);
    }
    int deviation = Integer.MAX_VALUE;
    while (true) {
        int max = maxHeap.poll();
        deviation = Math.min(deviation, max - min);
        if (max % 2 == 1) {
            break;                           // the max cannot shrink any further
        }
        max /= 2;
        min = Math.min(min, max);            // the halved value may be the new minimum
        maxHeap.offer(max);
    }
    return deviation;
}

# python
# LC 767 - Reorganize String
# IDEA: feasible iff max_freq <= (n + 1) // 2; then pop the two most frequent letters each round
# time = O(n log 26) = O(n), space = O(26)
import heapq
from collections import Counter

def reorganizeString(s):
    count = Counter(s)
    if max(count.values()) > (len(s) + 1) // 2:
        return ""                        # the most frequent letter cannot be spread out
    heap = [(-cnt, ch) for ch, cnt in count.items()]
    heapq.heapify(heap)
    res = []
    while len(heap) >= 2:
        cnt1, ch1 = heapq.heappop(heap)  # the two most frequent letters, different by construction
        cnt2, ch2 = heapq.heappop(heap)
        res.append(ch1)
        res.append(ch2)
        if cnt1 + 1 < 0:
            heapq.heappush(heap, (cnt1 + 1, ch1))
        if cnt2 + 1 < 0:
            heapq.heappush(heap, (cnt2 + 1, ch2))
    if heap:
        res.append(heap[0][1])           # one letter with one copy left
    return "".join(res)

// java
// LC 767 - Reorganize String
// IDEA: feasible iff maxFreq <= (n + 1) / 2; then poll the two most frequent letters each round
// time = O(n log 26) = O(n), space = O(26)
public String reorganizeString(String s) {
    int[] freq = new int[26];
    for (char c : s.toCharArray()) {
        freq[c - 'a']++;
    }
    int n = s.length();
    for (int f : freq) {
        if (f > (n + 1) / 2) {
            return "";                       // the most frequent letter cannot be spread out
        }
    }
    PriorityQueue<int[]> maxHeap = new PriorityQueue<>((a, b) -> Integer.compare(b[0], a[0]));   // {count, letter}
    for (int i = 0; i < 26; i++) {
        if (freq[i] > 0) {
            maxHeap.offer(new int[]{freq[i], i});
        }
    }
    StringBuilder sb = new StringBuilder();
    while (maxHeap.size() >= 2) {
        int[] first = maxHeap.poll();        // the two most frequent, different letters
        int[] second = maxHeap.poll();
        sb.append((char) (first[1] + 'a'));
        sb.append((char) (second[1] + 'a'));
        if (--first[0] > 0) maxHeap.offer(first);
        if (--second[0] > 0) maxHeap.offer(second);
    }
    if (!maxHeap.isEmpty()) {
        sb.append((char) (maxHeap.poll()[1] + 'a'));   // one letter with one copy left
    }
    return sb.toString();
}

# python
# LC 767 - Reorganize String (even/odd fill)
# IDEA: most frequent letter first into the even slots, then wrap to the odd ones
# time = O(n), space = O(n)
from collections import Counter

def reorganizeString_fill(s):
    count = Counter(s)
    if max(count.values()) > (len(s) + 1) // 2:
        return ""
    res = [''] * len(s)
    i = 0
    for ch, cnt in count.most_common():  # most frequent first: it must fit the even slots
        for _ in range(cnt):
            if i >= len(s):
                i = 1                    # even slots used up: continue on the odd ones
            res[i] = ch
            i += 2
    return "".join(res)

tasks = AAABBB, n = 2   max_freq = 3, tied for it = 2 (A and B)

  | A B idle | A B idle | A B
  \_ n + 1 _/\_ n + 1 _/\_ 2 _/

  slots = (3 - 1) * (2 + 1) + 2 = 8

# python
# LC 621 - Task Scheduler
# IDEA: (max_freq - 1) frames of width n + 1, plus every task tied for max_freq; never below len(tasks)
# time = O(N), space = O(26)
from collections import Counter

def leastInterval(tasks, n):
    count = Counter(tasks)
    max_freq = max(count.values())
    num_max = sum(1 for f in count.values() if f == max_freq)   # tasks tied for the max
    frames = (max_freq - 1) * (n + 1) + num_max
    return max(len(tasks), frames)       # enough distinct tasks fill every gap: no idle

// java
// LC 621 - Task Scheduler
// IDEA: (maxFreq - 1) frames of width n + 1, plus every task tied for maxFreq; never below tasks.length
// time = O(N), space = O(26)
public int leastInterval(char[] tasks, int n) {
    int[] freq = new int[26];
    int maxFreq = 0;
    for (char t : tasks) {
        freq[t - 'A']++;
        maxFreq = Math.max(maxFreq, freq[t - 'A']);
    }
    int numMax = 0;                          // tasks tied for the max
    for (int f : freq) {
        if (f == maxFreq) numMax++;
    }
    int frames = (maxFreq - 1) * (n + 1) + numMax;
    return Math.max(tasks.length, frames);   // enough distinct tasks fill every gap: no idle
}

# python
# LC 621 - Task Scheduler (simulation)
# IDEA: max-heap of remaining counts; a run task waits in a FIFO until time + n, then returns
# time = O(N * n) in the worst case (idle ticks are simulated), space = O(26)
import heapq
from collections import Counter, deque

def leastInterval_heap(tasks, n):
    heap = [-cnt for cnt in Counter(tasks).values()]   # most remaining on top
    heapq.heapify(heap)
    cooldown = deque()                   # (negated remaining, time it may run again)
    time = 0
    while heap or cooldown:
        time += 1
        if heap:
            remaining = heapq.heappop(heap) + 1        # run one copy
            if remaining < 0:
                cooldown.append((remaining, time + n))
        if cooldown and cooldown[0][1] == time:
            heapq.heappush(heap, cooldown.popleft()[0])
    return time