2020-01-26 432. All O`one Data Structure

ruicore · ruicore · commit bacd9eef5268 · 2020-01-26T16:37:26.000+08:00
diff --git a/LeetCode/2020-01-26-432-All-Oone-Data-Structure.py b/LeetCode/2020-01-26-432-All-Oone-Data-Structure.py
@@ -0,0 +1,87 @@
+# -*- coding: utf-8 -*-
+# @Author:             何睿
+# @Create Date:        2020-01-26 16:07:05
+# @Last Modified by:   何睿
+# @Last Modified time: 2020-01-26 16:35:40
+
+
+class Block(object):
+    def __init__(self, val=0):
+        self.val = val
+        self.keys = set()
+        self.before = None
+        self.after = None
+
+    def remove(self):
+        self.before.after = self.after
+        self.after.before = self.before
+        self.before, self.after = None, None
+
+    def insert_after(self, new_block):
+        old_after = self.after
+        self.after = new_block
+        new_block.before = self
+        new_block.after = old_after
+        old_after.before = new_block
+
+
+class AllOne(object):
+    def __init__(self):
+        self.begin = Block()  # sentinel
+        self.end = Block()  # sentinel
+        self.begin.after = self.end
+        self.end.before = self.begin
+        self.mapping = {}  # key to block
+
+    def inc(self, key):
+        if not key in self.mapping:  # find current block and remove key
+            current_block = self.begin
+        else:
+            current_block = self.mapping[key]
+            current_block.keys.remove(key)
+
+        if current_block.val + 1 != current_block.after.val:  # insert new block
+            new_block = Block(current_block.val + 1)
+            current_block.insert_after(new_block)
+        else:
+            new_block = current_block.after
+
+        new_block.keys.add(key)  # update new_block
+        self.mapping[key] = new_block  # ... and mapping of key to new_block
+
+        if not current_block.keys and current_block.val != 0:  # delete current block if not seninel
+            current_block.remove()
+
+    def dec(self, key):
+        if not key in self.mapping:
+            return
+
+        current_block = self.mapping[key]
+        del self.mapping[key]  # could use self.mapping.pop(key)
+        current_block.keys.remove(key)
+
+        if current_block.val != 1:
+            if current_block.val - 1 != current_block.before.val:  # insert new block
+                new_block = Block(current_block.val - 1)
+                current_block.before.insert_after(new_block)
+            else:
+                new_block = current_block.before
+            new_block.keys.add(key)
+            self.mapping[key] = new_block
+
+        if not current_block.keys:  # delete current block
+            current_block.remove()
+
+    def getMaxKey(self):
+        if self.end.before.val == 0:
+            return ""
+        key = self.end.before.keys.pop()  # pop and add back to get arbitrary (but not random) element
+        self.end.before.keys.add(key)
+        return key
+
+    def getMinKey(self):
+        if self.begin.after.val == 0:
+            return ""
+        key = self.begin.after.keys.pop()
+        self.begin.after.keys.add(key)
+        return key
