BinarySearch 0767 Frequency Stack

Problem statement

https://binarysearch.com/problems/Frequency-Stack/

Solution

Equal to Leetcode 0895. Maximum Frequency Stack.

Complexity

Time is O(1) for both operations, space is O(n).

Code

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

    def append(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