895. Maximum Frequency Stack

[zeikar]

Problem link

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

Problem Summary

많이 나온 원소부터 pop이 되는 스택을 구현하는 것이다.

Solution

하드 치고는 꽤 쉬운 문제.
문제 그대로 많이 나온 원소부터 pop이 되도록 구현하면 되는데 원소가 나온 카운트에 대한 맵을 하나 두고 스택도 카운트 별로 만들어 둔다. 예제에서 push, pop 쿼리가 최대 2만개 까지라고 했으므로 20001개의 스택을 만들어 두었다.

push의 경우 원소의 개수에 대한 스택에 집어넣으면 되고 최대 카운트를 갱신해준다. pop 할 때는 최대 카운트의 스택에서 하나씩 빼주고 없다면 최대 카운트를 1 줄여주면 된다.

Source Code

import collections


class FreqStack:

    def __init__(self):
        self.count = collections.defaultdict(int)
        self.stack = [collections.deque() for _ in range(20000)]
        self.max_freq = 0

    def push(self, val: int) -> None:
        count = self.count[val] if val in self.count else 0
        self.count[val] = count + 1

        if count + 1 > self.max_freq:
            self.max_freq = count + 1

        self.stack[count + 1].append(val)

    def pop(self) -> int:
        val = self.stack[self.max_freq].pop()
        self.count[val] -= 1
        if not self.stack[self.max_freq]:
            self.max_freq -= 1
        return val

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