LeetCode-Solutions/Python/Maximum Frequency Stack.py at main · ToDevelopersTeam/LeetCode-Solutions · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
from collections import defaultdict
from heapq import heappop, heappush


class FreqStack:
    """
    Implementation of a frequency stack based on a stack of stacks

    Memory:
        creation - O(n)
        'push' and 'pop' - O(1)

    Time:
        creation - O(n)
        'push' and 'pop' - O(1)
    """

    def __init__(self):
        self.frequency = defaultdict(int)
        self.stack = defaultdict(list)
        self.max_freq = 0

    def push(self, val: int) -> None:
        self.frequency[val] += 1
        self.max_freq = max(self.frequency[val], self.max_freq)
        self.stack[self.frequency[val]].append(val)

    def pop(self) -> int:
        val = self.stack[self.max_freq].pop()
        self.frequency[val] -= 1

        if not self.stack[self.max_freq]:
            self.max_freq -= 1

        return val


class FreqStack:
    """
    Implementation of a frequency stack based on a heap

    Memory:
        creation - O(n)
        'push' and 'pop' - O(1)

    Time:
        creation - O(n)
        'push' and 'pop' - O(log(n))
    """

    def __init__(self):
        self.frequency = defaultdict(int)
        self.heap = []
        self.index = 0

    def push(self, val: int) -> None:
        self.frequency[val] += 1
        self.index += 1
        heappush(self.heap, (-self.frequency[val], -self.index, val))

    def pop(self) -> int:
        _, _, val = heappop(self.heap)
        self.frequency[val] -= 1
        return val