LeetCode 풀기 - 895. Maximum Frequency Stack

895. Maximum Frequency Stack

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

Maximum Frequency Stack - LeetCode

문제)

솔루션1) 시간 초과 발생

시간 초과가 발생합니다.

pop()을 수행할 때마다 sorted()로 정렬을 수행하기 때문에 너무 느립니다.

push() 수행할 때 이미 정렬을 수행해 놓고 pop()에서는 O(1)으로 가져오는 방식으로 최적화가 필요합니다.

from collections import defaultdict
class FreqStack:
    def __init__(self):
        self.stack = []
        self.counter = defaultdict(int)

    def push(self, val):
        self.stack.append(val)
        self.counter[val] += 1

    def pop(self):
        sorted_counter = sorted(self.counter.items(), key=lambda x: x[1], reverse=True)
        most_freq = sorted_counter[0][1]
        most_freq_vals = [x[0] for x in list(filter(lambda x: x[1] == most_freq, sorted_counter))]

        for i in range(len(self.stack)-1, -1, -1):
            if self.stack[i] in most_freq_vals:
                self.counter[self.stack[i]] -= 1
                ret = self.stack.pop(i)
                break
        return ret

솔루션2) heapq 사용

heapq를 사용해서 push()가 호출될 때마다 heapq에 push를 수행합니다.

frequency가 크고 최근에 추가된 것으로 정렬되도록 우선수위를 조정합니다.

import heapq as hq
from collections import defaultdict

class FreqStack:
    def __init__(self):
        self.freq = defaultdict(int)
        self.counter = 0
        self.heap = []

    def push(self, val):
        self.freq[val] += 1

        # frequency가 크고 나중에 추가될 수록 우선순위가 높음
        hq.heappush(self.heap, (-self.freq[val], -self.counter, val))

        self.counter += 1

    def pop(self):
        val = hq.heappop(self.heap)[2]
        self.freq[val] -= 1
        return val