Dinner Plate Stacks

Problem Id: 1172 Difficulty: Hard


Intuition

Solution


#!/usr/bin/env python
# coding=utf-8
#
# Author: Lucas
# Date: 2019-08-25 12:02:40


import heapq


class DinnerPlates:
    def __init__(self, capacity: int):
        self.stacks = []
        self.cap = capacity
        self.empty_prev = []

    def _rm_empty_stack(self):
        while self.stacks and not self.stacks[-1]:
            self.stacks.pop()

    def push(self, val: int) -> None:
        if self.empty_prev and self.empty_prev[0] >= len(self.stacks):
            self.empty_prev = []
        if not self.stacks:
            self.stacks.append([])

        if self.empty_prev:
            i = heapq.heappop(self.empty_prev)
            self.stacks[i].append(val)
        else:
            if len(self.stacks[-1]) == self.cap:
                self.stacks.append([])
            self.stacks[-1].append(val)

    def pop(self) -> int:
        if not self.stacks:
            return -1
        ans = self.stacks[-1].pop()
        while self.stacks and not self.stacks[-1]:
            self.stacks.pop()
        return ans

    def popAtStack(self, index: int) -> int:
        if index >= len(self.stacks) or not self.stacks[index]:
            return -1
        heapq.heappush(self.empty_prev, index)
        ans = self.stacks[index].pop()
        while self.stacks and not self.stacks[-1]:
            self.stacks.pop()
        return ans


# Your DinnerPlates object will be instantiated and called as such:
# obj = DinnerPlates(capacity)
# obj.push(val)
# param_2 = obj.pop()
# param_3 = obj.popAtStack(index)