https://leetcode.com/problems/maximum-frequency-stack

Solution 1, using SortedList

There is solution with complexity O( log n) for pop and push. We will keep 3 pieces of information:

  1. self.SList is sorted list, where we keep tuples of 3 numbers: frequency, number itself and biggest index we have for this number.
  2. self.Places is defaultdict, where for each num we keep all indexes.
  3. self.N number of element we push into our data structure. Note, that we do not decrease this number, only increase.

Now, how function will work:

  1. push(self, x): first, we check frequency of new element and update it: we remove tuple from our sorted list and put updated frequency. We also add new index to self.Places[x] and update self.N
  2. pop(self): we pop last element from our sorted list: it will be what we wanted. If frequency of extracted element is 2 or more, than we need to put element with updated frequency to our sorted list. Also we remove last element from self.Places.

Complexity: all operations which we use when we work with sorted list have complexity O(log n), so final complexity of both push and pop is just O(log n)

from sortedcontainers import SortedList

class FreqStack:
    def __init__(self):
        self.SList = SortedList()   #(freq, num, last)
        self.Places = defaultdict(list)  #num -> all indexes
        self.N = 0

    def push(self, x):
        freq_x = len(self.Places[x])
        if freq_x > 0: self.SList.remove((freq_x, self.Places[x][-1], x))
        self.SList.add((freq_x + 1, self.N, x))
        self.Places[x].append(self.N)
        self.N += 1
        
    def pop(self):
        freq, end, num = self.SList.pop()
        self.Places[num].pop()
        if freq > 1: self.SList.add((freq - 1, self.Places[num][-1], num))
        return num

Solution 2: using stacks

There is also O(1) pop/push solution: let us look at it at the following example: we push into stack numbers[0, 7, 3, 7, 8, 8, 8, 3, 4, 0]. Then we will keep dictionary of stacks: {1: [0, 7, 3, 8, 4], 2: [7, 8, 3, 0], 3: [8]}. What is logic behind this? We have 3 items to keep:

  1. self.freq: frequencies of each element
  2. self.stacks: stack for each frequency, note that it is not stack for each element with indexes, but something different: for each frequency we will keep elements we have with this frequency in the order like in original stack. self.maxfreq is maximal frequency.

Now, how functions will work:

  1. push(self, x): when we add new element, first we need to increase it frequency. If new frequency is more than maximum frequency we have, we update it. Finally, we update our self.stacks for new element: put it to the end of stack.
  2. pop(self): first of all, element we are looking for is last element in self.stacks[self.maxfreq], so we pop it. Also we need to decrease frequency of this element. Finally, if number of elements with highest frequency equal to 0, we can remove this stack. Then we return x.

Complexity: each operation we do have just O(1) complexity, so both push and pop work just in O(1).

class FreqStack:
    def __init__(self):
        self.freq = Counter()
        self.stacks = defaultdict(list)
        self.maxfreq = 0

    def push(self, x):
        freq_x = self.freq[x] + 1
        self.freq[x] = freq_x
        self.maxfreq = max(self.maxfreq, freq_x)
        self.stacks[freq_x].append(x)

    def pop(self):
        x = self.stacks[self.maxfreq].pop()
        self.freq[x] -= 1
        if not self.stacks[self.maxfreq]: self.maxfreq -= 1
        return x

If you like the solution, you can upvote it on leetcode discussion section: Problem 0895