Maximum Frequency Stack

hardAmazon-tagged: Design-style coding46 stepsJava · Python · C++

pop removes the most frequent element, breaking ties by recency. Keep each value's frequency and a stack for every frequency level. Pushing x with new frequency f puts it on stack f; popping takes from the highest non-empty level.

▶ Watch the animated trace

Traced input: push(5), push(7), push(5), push(7), push(4), push(5), pop(), pop(), pop(), pop() → [null, null, null, null, null, null, null, 5, 7, 5, 4]. Original problem: neetcode.io/problems/maximum-frequency-stack

Complexity

BestO(1)

AverageO(1)

WorstO(1)

SpaceO(n)

Both operations are O(1).

Key idea

An element with frequency 3 sits on stacks 1, 2 and 3: one copy per level it reached.
The top of the highest stack is the most frequent and, among ties, the most recent.
After popping, decrease its frequency; if that stack empties, lower maxFreq.

Common mistakes

A heap of (freq, time) works but is O(log n).
Forgetting to lower maxFreq when its stack becomes empty.

Solution

Java

class FreqStack {
    private final Map<Integer, Integer> freq = new HashMap<>();
    private final Map<Integer, Deque<Integer>> group = new HashMap<>();
    private int maxFreq = 0;

    void push(int x) {
        int f = freq.merge(x, 1, Integer::sum);
        maxFreq = Math.max(maxFreq, f);
        group.computeIfAbsent(f, k -> new ArrayDeque<>()).push(x);
    }

    int pop() {
        int x = group.get(maxFreq).pop();
        freq.merge(x, -1, Integer::sum);
        if (group.get(maxFreq).isEmpty()) maxFreq--;
        return x;
    }
}

Python

class FreqStack:
    def __init__(self):
        self.freq = Counter()
        self.group = defaultdict(list)
        self.max_freq = 0

    def push(self, x: int) -> None:
        self.freq[x] += 1
        f = self.freq[x]
        self.max_freq = max(self.max_freq, f)
        self.group[f].append(x)

    def pop(self) -> int:
        x = self.group[self.max_freq].pop()
        self.freq[x] -= 1
        if not self.group[self.max_freq]: self.max_freq -= 1
        return x

C++

class FreqStack {
    unordered_map<int, int> freq;
    unordered_map<int, vector<int>> group;
    int maxFreq = 0;

public:
    void push(int x) {
        int f = ++freq[x];
        maxFreq = max(maxFreq, f);
        group[f].push_back(x);
    }

    int pop() {
        int x = group[maxFreq].back();
        group[maxFreq].pop_back();
        freq[x]--;
        if (group[maxFreq].empty()) maxFreq--;
        return x;
    }
};

Edge cases

Each one has its own trace in the interactive view.

Case	Input	Answer
All different All frequency 1: behaves like a normal stack.	`push(1), push(2), push(3), pop(), pop()`	`[null, null, null, null, 3, 2]`
Single value maxFreq drops from 2 to 1 between pops.	`push(9), push(9), pop(), pop()`	`[null, null, null, 9, 9]`