maximum-frequency-stack.py

import heapq
from collections import Counter, defaultdict


class FreqStack:

    def __init__(self):
        self._counts = Counter()
        self._freq_stack = defaultdict(list)
        self._max_freq = 0

    def push(self, x: int) -> None:
        self._counts[x] += 1
        self._freq_stack[self._counts[x]].append(x)
        self._max_freq = max(self._counts[x], self._max_freq)

    def pop(self) -> int:
        val = self._freq_stack[self._max_freq].pop()
        self._counts[val] -= 1

        if not self._freq_stack[self._max_freq]:
            self._max_freq -= 1

        return val


class FreqStackHeap:  # log N push and pop, suboptimal

    def __init__(self):
        self._counts = Counter()
        self._heap = []
        self._pos = 0

    def push(self, x: int) -> None:
        self._counts[x] += 1
        self._pos += 1  # should be mod maxint in the real code
        heapq.heappush(self._heap, (-self._counts[x], -(self._pos + 1), x))

    def pop(self) -> int:
        _, _, val = heapq.heappop(self._heap)
        self._counts[val] -= 1
        return val



# Your FreqStack object will be instantiated and called as such:
# obj = FreqStack()
# obj.push(x)
# param_2 = obj.pop()