Problem statement

https://leetcode.com/problems/dinner-plate-stacks/

Solution 1

The idea is to keep stack of plates self.stack and self.S is sorted list of pairs (ind, size) of not-full plates.

Then when we push we need to check:

  1. If we have empty self.S, it means that all stack are full, so we need to create new one, add it.
  2. If it is not empty, take the smallest element from self.S: pair (ind, sz).
  3. If after adding new element to this stack it is not full, we update information in self.S.
  4. Finally, add this plate to self.plates.

When we want to perform pop, we:

  1. First, remove all empty stacks from the end and if we do not have anything in the end, we return -1.
  2. Now, we know that the stack we want to work with is the last on in self.plates, so we update info.

Update is auxiliary function, which given index and size of stack from which we pop, update self.plates and self.S: if capacity was not full, we remove old information from self.S (if it was full, it was not there) and then we add new info. Also update self.plates.

Finally, popAtStack function will work like this:

  1. If ind is more than length of self.plates or stack is empty, return -1.
  2. If it is not empty, use update function.

Complexity

Time complexity of all operations is O(log n).

Code

from sortedcontainers import SortedList

class DinnerPlates:
    def __init__(self, capacity):
        self.cap = capacity
        self.plates = []
        self.S = SortedList()

    def push(self, val):
        if len(self.S) == 0:
            ind, sz = len(self.plates), 0
            self.plates.append([])
        else:
            ind, sz = self.S.pop(0)
            
        if sz + 1 < self.cap:
            self.S.add((ind, sz + 1))
            
        self.plates[ind].append(val)
        
    def update(self, ind, sz):
        if self.cap != sz: self.S.remove((ind, sz))
        self.S.add((ind, sz - 1))
        return self.plates[ind].pop()
        
    def pop(self):
        while self.plates and len(self.plates[-1]) == 0:
            self.S.pop(-1)
            self.plates.pop()
        
        if not self.plates: return -1
        
        ind, sz = len(self.plates) - 1, len(self.plates[-1])
        return self.update(ind, sz)

    def popAtStack(self, ind):
        if len(self.plates) < ind: return -1
        sz = len(self.plates[ind])
        
        if sz == 0: return -1
        return self.update(ind, sz)

Solution 2

There is better approach using heaps: idea is similar, but now what we keep in heap is places which are available. Also we can have several equal elements in heap.

  1. When we push new element and if we do not have places, we put val and create self.c - 1 empty places. If we have places, we check if they are in existing stack, if they are not (it can happen that we made some stack fully empty) in this case we need to start new stack and add val.
  2. When we pop element first we remove last empty stacks. Then we add one more vacant place in this stack and return popped element.
  3. Finally, popAtStack: we check if our index is more than length of plates or stack in empty, in this case we return -1. Then we extract res, add one more vacant place and return res.

Complexity

Time complexity for push is O(log n) for other two is O(1).

Code

class DinnerPlates:
    def __init__(self, capacity):
        self.plates = []
        self.c = capacity
        self.emp = []

    def push(self, val):
        if not self.emp:
            self.plates += [[val]]
            for _ in range(self.c - 1): heappush(self.emp, len(self.plates)-1)
            return
            
        ind = heappop(self.emp)
        if ind == len(self.plates): self.plates += [[]]
        self.plates[ind] += [val]
       

    def pop(self):
        while len(self.plates) > 0 and not self.plates[-1]:
            self.plates.pop()
            
        if not self.plates: return -1
        heapq.heappush(self.emp, len(self.plates) - 1)
        return self.plates[-1].pop()

    def popAtStack(self, ind):
        if ind >= len(self.plates) or not self.plates[ind]: return -1
        heapq.heappush(self.emp, ind)
        return self.plates[ind].pop()