All O`one Data Structure LeetCode solution in Python

1
from collections import defaultdict
2

3

4
class Set(object):
5

6
    def __init__(self):
7
        self._dict = {}
8
        self._list = []
9
        self._len = 0
10

11
    def add(self, v):
12
        if v in self._dict:
13
            pass
14
        else:
15
            self._list.append(v)
16
            self._dict[v] = self._len
17
            self._len += 1
18

19
    def __contains__(self, item):
20
        return item in self._dict
21

22
    def remove(self, v):
23
        if v not in self._dict:
24
            pass
25
        else:
26
            idx = self._dict[v]
27

28
            if idx < self._len - 1:
29
                self._list[idx] = self._list[-1]
30
                self._dict[self._list[idx]] = idx
31

32
            del self._dict[v]
33
            self._list.pop()
34
            self._len -= 1
35

36
    def __iter__(self):
37
        return iter(self._list)
38

39
    def check(self):
40
        assert len(self._dict) == len(self._list)
41
        for idx, key in enumerate(self._list):
42
            assert self._dict[key] == idx
43

44
    def __len__(self):
45
        return self._len
46

47
    def __repr__(self):
48
        return f"{self._dict}\n{self._list}"
49

50

51
class Node(object):
52

53
    def __init__(self, value):
54
        self.value = value
55
        self.prev = None
56
        self.next = None
57

58
    def set_prev(self, other_node):
59
        self.prev = other_node
60
        if other_node is not None:
61
            other_node.next = self
62

63
    def set_next(self, other_node):
64
        self.next = other_node
65
        if other_node is not None:
66
            other_node.prev = self
67

68
    def __repr__(self):
69
        return str(self.value)
70

71

72
class SortedKeyList(object):
73

74
    def __init__(self):
75
        self.head = Node(value=0)
76
        self.tail = Node(value=float("inf"))
77
        self.head.set_next(self.tail)
78

79
    def is_empty(self):
80
        return self.head.next is self.tail
81

82
    def min_key(self):
83
        if self.is_empty():
84
            return None
85
        else:
86
            return self.head.next.value
87

88
    def max_key(self):
89
        if self.is_empty():
90
            return None
91
        else:
92
            return self.tail.prev.value
93

94
    def incr_key(self, orig_node):
95
        orig_value = orig_node.value
96
        new_value = orig_value + 1
97
        next_node = orig_node.next
98
        if next_node.value == new_value:
99
            return self
100
        else:
101
            new_node = Node(new_value)
102
            new_node.set_prev(orig_node)
103
            new_node.set_next(next_node)
104
            return self
105

106
    def decr_key(self, orig_node):
107
        orig_value = orig_node.value
108
        new_value = orig_value - 1
109
        prev_node = orig_node.prev
110
        if prev_node.value == new_value:
111
            return self
112
        else:
113
            new_node = Node(new_value)
114
            new_node.set_next(orig_node)
115
            new_node.set_prev(prev_node)
116
            return self
117

118
    def delete_node(self, node):
119
        prev_node = node.prev
120
        next_node = node.next
121
        prev_node.set_next(next_node)
122
        return self
123

124
    def __repr__(self):
125
        a = []
126
        node = self.head.next
127
        while True:
128
            if node is self.tail:
129
                break
130

131
            assert node.next.prev is node
132
            assert node.prev.next is node
133

134
            a.append(node.value)
135
            node = node.next
136
        return str(a)
137

138

139
class AllOne:
140

141
    def __init__(self):
142
        self.count_list = SortedKeyList()
143
        self.counter = defaultdict(lambda: self.count_list.head)
144
        self.keyed_by_count = defaultdict(set)
145

146
    def __repr__(self):
147
        return f"count_list={self.count_list}\ncounter={dict(self.counter)}\nkeyed_by_count={dict(self.keyed_by_count)}"
148

149
    def inc(self, key: str) -> None:
150
        orig_count_node = self.counter[key]
151
        orig_count = orig_count_node.value
152

153
        self.count_list.incr_key(orig_count_node)
154
        new_count_node = orig_count_node.next
155
        new_count = new_count_node.value
156
        assert new_count == orig_count + 1
157

158
        self.counter[key] = new_count_node
159

160
        if key in self.keyed_by_count[orig_count]:
161
            self.keyed_by_count[orig_count].remove(key)
162
        self.keyed_by_count[new_count].add(key)
163

164
        if len(self.keyed_by_count[orig_count]) == 0:
165
            del self.keyed_by_count[orig_count]
166
            if orig_count_node.value > 0:
167
                self.count_list.delete_node(orig_count_node)
168

169
    def dec(self, key: str) -> None:
170
        orig_count_node = self.counter[key]
171
        orig_count = orig_count_node.value
172

173
        self.count_list.decr_key(orig_count_node)
174
        new_count_node = orig_count_node.prev
175
        new_count = new_count_node.value
176
        assert new_count == orig_count - 1
177

178
        self.counter[key] = new_count_node
179

180
        if key in self.keyed_by_count[orig_count]:
181
            self.keyed_by_count[orig_count].remove(key)
182
        self.keyed_by_count[new_count].add(key)
183

184
        if new_count == 0:
185
            del self.counter[key]
186

187
        if len(self.keyed_by_count[orig_count]) == 0:
188
            del self.keyed_by_count[orig_count]
189
            self.count_list.delete_node(orig_count_node)
190

191
    def getMaxKey(self) -> str:
192
        max_count = self.count_list.max_key()
193
        if max_count is not None:
194
            return next(iter(self.keyed_by_count[max_count]))
195
        else:
196
            return ""
197

198
    def getMinKey(self) -> str:
199
        min_count = self.count_list.min_key()
200
        if min_count is not None:
201
            return next(iter(self.keyed_by_count[min_count]))
202
        else:
203
            return ""
204

205

206
def drive(m, p):
207
    s = AllOne()
208

209
    naive_counter = defaultdict(int)
210

211
    for method, param in zip(m, p):
212
        if method in ("inc", "dec"):
213
            word = param[0]
214

215
            if method == "inc":
216
                s.inc(word)
217
                naive_counter[word] += 1
218

219
            elif method == "dec":
220
                s.dec(word)
221
                naive_counter[word] -= 1
222

223
            if naive_counter[word] == 0:
224
                del naive_counter[word]
225

226
            tmp_counter = defaultdict(set)
227
            for k, v in naive_counter.items():
228
                tmp_counter[v].add(k)
229

230
            sorted_keys = sorted(tmp_counter.keys())
231
            min_key = sorted_keys[0]
232
            max_key = sorted_keys[-1]
233

234
            if s.getMaxKey() not in tmp_counter[max_key]:
235
                print("Oh No!!!")
236
                return s, naive_counter, tmp_counter
237

238
            if s.getMinKey() not in tmp_counter[min_key]:
239
                print("Oh No!!!")
240
                return s, naive_counter, tmp_counter
241

242
    return None, None, None
243

244

245
if __name__ == "__main__":
246
    m = [
247
        "AllOne",
248
        "inc",
249
        "inc",
250
        "getMaxKey",
251
        "getMinKey",
252
        "inc",
253
        "getMaxKey",
254
        "getMinKey",
255
    ]
256
    p = [[], ["hello"], ["hello"], [], [], ["leet"], [], []]
257
    s, naive_counter, tmp_counter = drive(m, p)