All O`one Data Structure

Problem Id: 432 Difficulty: Hard


Intuition

Solution


#!/usr/bin/env python
# coding=utf-8
#
# Author: Lucas
# Date: 2019-07-13 23:51:51


class AllOne:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.counts = {}
        self.keys = {}
        self.min_key = 1
        self.max_key = 0

    def inc(self, key: str) -> None:
        """
        Inserts a new key <Key> with value 1. Or increments an existing key by 1.
        """
        if key not in self.counts:
            self.counts[key] = 1
        else:
            self.counts[key] += 1

        prev = self.counts[key] - 1
        if prev:
            self.keys[prev].remove(key)
            if not self.keys[prev]:
                self.keys.pop(prev)
        count = prev + 1
        if count not in self.keys:
            self.keys[count] = set()
        self.keys[count].add(key)

        if count < self.min_key or (prev == self.min_key and prev not in self.keys):
            self.min_key = count

        if count > self.max_key:
            self.max_key = count

    def dec(self, key: str) -> None:
        """
        Decrements an existing key by 1. If Key's value is 1, remove it from the data structure.
        """
        if key not in self.counts:
            return

        self.counts[key] -= 1
        if self.counts[key] == 0:
            self.counts.pop(key)

        count = self.counts[key] if key in self.counts else 0
        prev = count + 1
        self.keys[prev].remove(key)
        if not self.keys[prev]:
            self.keys.pop(prev)
        if count:
            if count not in self.keys:
                self.keys[count] = set()
            self.keys[count].add(key)

        if count:
            if count < self.min_key:
                self.min_key = count
        else:
            self.min_key = min(self.keys.keys())

        if prev == self.max_key and prev not in self.keys:
            self.max_key = count

    def getMaxKey(self) -> str:
        """
        Returns one of the keys with maximal value.
        """
        if not self.counts:
            return ''
        key = self.keys[self.max_key].pop()
        self.keys[self.max_key].add(key)
        return key

    def getMinKey(self) -> str:
        """
        Returns one of the keys with Minimal value.
        """
        if not self.counts:
            return ''
        key = self.keys[self.min_key].pop()
        self.keys[self.min_key].add(key)
        return key


# Your AllOne object will be instantiated and called as such:
# obj = AllOne()
# obj.inc(key)
# obj.dec(key)
# param_3 = obj.getMaxKey()
# param_4 = obj.getMinKey()