895. Maximum Frequency Stack

Problem¶

Support pushing integers and popping the most frequent value; break ties by its most recent remaining push. Pops are guaranteed to occur on a nonempty stack.

push: 5, 7, 5, 7, 4, 5
pop:  5, 7, 5, 4
remaining stack, bottom to top: [5, 7]

class FreqStack:
    def __init__(self):
        pass

    def push(self, val: int) -> None:
        pass

    def pop(self) -> int:
        pass

Key trick¶

Maintain a stack for each frequency: every push records the value at its new frequency, and every pop uses the highest-frequency stack.

Trap¶

Lower-frequency entries are historical records, not stale entries: preserve them, and do not append the popped value to a lower-frequency stack.

Why is it interesting?¶

Grouping by frequency turns a priority-ordering problem into constant-time stack operations without a heap or timestamps.

Python solution¶

class FreqStack:
    def __init__(self):
        self.freqs = {}
        self.buckets = [[]]  # Index is frequency; index 0 is unused.

    def push(self, val: int) -> None:
        freq = self.freqs.get(val, 0) + 1
        self.freqs[val] = freq

        if freq == len(self.buckets):
            self.buckets.append([])

        # Arrival order at this frequency resolves recency ties.
        self.buckets[freq].append(val)

    def pop(self) -> int:
        val = self.buckets[-1].pop()

        self.freqs[val] -= 1
        if self.freqs[val] == 0:
            del self.freqs[val]

        if not self.buckets[-1]:
            self.buckets.pop()

        # Its earlier occurrence already exists in the lower group.
        return val

Time: Expected amortized \(O(1)\) per operation.

Space: \(O(n)\) for \(n\) currently stored elements.

Comment on my solution¶

Your indexed max-heap approach is viable, but the implementation has three correctness bugs:

New values never sift upward. Pushing distinct values must promote the newest one because their frequencies tie.
_swap updates the wrong dictionary keys. Use heap[i][3] and heap[j][3], the values, rather than heap[i][2][-1], the timestamps.
An exhausted root does not necessarily sink to the last array position. Remove it by replacing the root with the last entry, updating its index, and sifting down; checking only heap[-1] can leave zero-frequency entries behind.

Also, a pop can remove the last pushed element; you do not need a linked list because removing an arbitrary chronological position is unnecessary.

A corrected heap solution takes \(O(\log d)\) per operation for \(d\) distinct stored values, whereas frequency stacks achieve expected amortized \(O(1)\).

# Ideas:
# - It looks like problem Max stack which I just did.
#   - This use a doubly linked list for the stack
#   - and a max heap (node.idx is index of node in heap)
#   - node have a push_id to break heap ties
# - here we don't need a prev/next on node and a doubly linked list
#   because we never remove the last pushed node.
#   - push order only serve for tie breaking in the heap

# WRONG
class FreqStack:
    def __init__(self):
        self.heap = [] # heap of (frequency, push_id, [push_ids], val)
        self.values = {} # map value (not frequency for that value) to index in heap
        self.next_push_id = 0

    def push(self, val: int) -> None:
        push_id = self.next_push_id
        self.next_push_id += 1

        if val not in self.values:
            self.values[val] = len(self.heap)
            self.heap.append([1, push_id, [push_id], val])
            return

        i = self.values[val]
        item = self.heap[i]
        item[0] += 1 # increment freq
        item[1] = push_id
        item[2].append(push_id)
        self._update(i)

    def _swap(self, i, j):
        heap = self.heap
        heap[i], heap[j] = heap[j], heap[i]
        val_i = heap[i][2][-1]
        val_j = heap[j][2][-1]
        self.values[val_i] = i
        self.values[val_j] = j

    def _up(self, i):
        heap = self.heap
        while i:
            j = (i - 1) // 2
            if heap[j] >= heap[i]:
                break
            self._swap(i,j)
            i = j

    def _down(self, i):
        heap = self.heap
        n = len(heap)
        while 2 * i + 1 < n:
            j = 2 * i + 1
            if j + 1 < n and heap[j + 1] > heap[j]:
                j += 1
            if heap[i] >= heap[j]:
                break
            self._swap(i,j)
            i = j

    def _update(self, i):
        j = (i - 1) // 2
        if j >= 0 and self.heap[j] < self.heap[i]:
            self._up(i)
        else:
            self._down(i)

    def pop(self) -> int:
        item = self.heap[0]
        val = item[-1]
        item[0] -= 1 # decrement frequence
        item[2].pop() # remove the last pushed id
        item[1] = item[2][-1] if item[2] else -1 # update last pushed id
        self._update(0)
        if self.heap[-1][0] == 0:
            self.heap.pop()
            del self.values[val]
        return val

# WORKS
# 2 corrections by myself after reading AI comments
# AI corrected pop function
class FreqStack:
    def __init__(self):
        self.heap = []  # Entries: [frequency, latest_push_id, push_ids, val]
        self.values = {}  # val -> heap index
        self.next_push_id = 0

    def push(self, val: int) -> None:
        push_id = self.next_push_id
        self.next_push_id += 1

        if val not in self.values:
            i = len(self.heap)
            self.values[val] = i
            self.heap.append([1, push_id, [push_id], val])
            self._up(i)  # FIX: insertion must preserve heap order.
            return

        i = self.values[val]
        item = self.heap[i]
        item[0] += 1
        item[1] = push_id
        item[2].append(push_id)
        self._update(i)

    def _swap(self, i, j):
        heap = self.heap
        heap[i], heap[j] = heap[j], heap[i]
        # FIX: map actual values, not their latest push IDs.
        self.values[heap[i][3]] = i
        self.values[heap[j][3]] = j

    def _up(self, i):
        heap = self.heap
        while i:
            j = (i - 1) // 2
            if heap[j] >= heap[i]:
                break
            self._swap(i, j)
            i = j

    def _down(self, i):
        heap = self.heap
        n = len(heap)
        while 2 * i + 1 < n:
            j = 2 * i + 1
            if j + 1 < n and heap[j + 1] > heap[j]:
                j += 1
            if heap[i] >= heap[j]:
                break
            self._swap(i, j)
            i = j

    def _update(self, i):
        j = (i - 1) // 2
        if j >= 0 and self.heap[j] < self.heap[i]:
            self._up(i)
        else:
            self._down(i)

    def pop(self) -> int:
        item = self.heap[0]
        val = item[-1]
        item[0] -= 1
        item[2].pop()
        item[1] = item[2][-1] if item[2] else -1

        if item[0] == 0:
            # FIX: move the exhausted entry to the end before deleting it.
            self._swap(0, len(self.heap) - 1)
            self.heap.pop()
            del self.values[val]

        if self.heap:
            self._update(0)

        return val