1172. Dinner Plate Stacks

You have an infinite number of stacks arranged in a row and numbered (left to right) from 0, each of the stacks has the same maximum capacity.

Implement the DinnerPlates class:

  • DinnerPlates(int capacity) Initializes the object with the maximum capacity of the stacks.
  • void push(int val) Pushes the given positive integer val into the leftmost stack with size less than capacity.
  • int pop() Returns the value at the top of the rightmost non-empty stack and removes it from that stack, and returns -1 if all stacks are empty.
  • int popAtStack(int index) Returns the value at the top of the stack with the given index and removes it from that stack, and returns -1 if the stack with that given index is empty.

Example:

Input: 
["DinnerPlates","push","push","push","push","push","popAtStack","push","push","popAtStack","popAtStack","pop","pop","pop","pop","pop"]
[[2],[1],[2],[3],[4],[5],[0],[20],[21],[0],[2],[],[],[],[],[]]
Output: 
[null,null,null,null,null,null,2,null,null,20,21,5,4,3,1,-1]

Explanation: 
DinnerPlates D = DinnerPlates(2);  // Initialize with capacity = 2
D.push(1);
D.push(2);
D.push(3);
D.push(4);
D.push(5);         // The stacks are now:  2  4
                                           1  3  5
                                           ﹈ ﹈ ﹈
D.popAtStack(0);   // Returns 2.  The stacks are now:     4
                                                       1  3  5
                                                       ﹈ ﹈ ﹈
D.push(20);        // The stacks are now: 20  4
                                           1  3  5
                                           ﹈ ﹈ ﹈
D.push(21);        // The stacks are now: 20  4 21
                                           1  3  5
                                           ﹈ ﹈ ﹈
D.popAtStack(0);   // Returns 20.  The stacks are now:     4 21
                                                        1  3  5
                                                        ﹈ ﹈ ﹈
D.popAtStack(2);   // Returns 21.  The stacks are now:     4
                                                        1  3  5
                                                        ﹈ ﹈ ﹈ 
D.pop()            // Returns 5.  The stacks are now:      4
                                                        1  3 
                                                        ﹈ ﹈  
D.pop()            // Returns 4.  The stacks are now:   1  3 
                                                        ﹈ ﹈   
D.pop()            // Returns 3.  The stacks are now:   1 
                                                        ﹈   
D.pop()            // Returns 1.  There are no stacks.
D.pop()            // Returns -1.  There are still no stacks.

 

Constraints:

  • 1 <= capacity <= 20000
  • 1 <= val <= 20000
  • 0 <= index <= 100000
  • At most 200000 calls will be made to push, pop, and popAtStack.

Rust Solution

use std::collections::HashMap;

#[derive(Debug)]
struct DinnerPlates {
    capacity: usize,
    stacks: HashMap<usize, Vec<i32>>,
    start: usize,
    end: usize,
}

impl DinnerPlates {
    fn new(capacity: i32) -> Self {
        let capacity = capacity as usize;
        let stacks = HashMap::new();
        let start = 0;
        let end = 0;
        DinnerPlates {
            capacity,
            stacks,
            start,
            end,
        }
    }

    fn push(&mut self, val: i32) {
        while self.stacks.entry(self.start).or_default().len() == self.capacity {
            self.start += 1;
        }
        self.stacks.entry(self.start).or_default().push(val);
        self.end = self.end.max(self.start);
    }

    fn pop(&mut self) -> i32 {
        while self.stacks.entry(self.end).or_default().is_empty() {
            if self.end > 0 {
                self.end -= 1;
            } else {
                return -1;
            }
        }
        self.pop_at_stack(self.end as i32)
    }

    fn pop_at_stack(&mut self, index: i32) -> i32 {
        let i = index as usize;
        if let Some(val) = self.stacks.entry(i).or_default().pop() {
            self.start = self.start.min(i);
            val
        } else {
            -1
        }
    }
}
#[test]
fn test() {
    let mut obj = DinnerPlates::new(2);
    obj.push(1);
    obj.push(2);
    obj.push(3);
    obj.push(4);
    obj.push(5);
    assert_eq!(obj.pop_at_stack(0), 2);
    obj.push(20);
    obj.push(21);
    assert_eq!(obj.pop_at_stack(0), 20);
    assert_eq!(obj.pop_at_stack(2), 21);
    assert_eq!(obj.pop(), 5);
    assert_eq!(obj.pop(), 4);
    assert_eq!(obj.pop(), 3);
    assert_eq!(obj.pop(), 1);
    assert_eq!(obj.pop(), -1);
    let mut obj = DinnerPlates::new(1);
    obj.push(1);
    obj.push(2);
    obj.push(3);
    assert_eq!(obj.pop_at_stack(1), 2);
    assert_eq!(obj.pop(), 3);
    assert_eq!(obj.pop(), 1);
}

Having problems with this solution? Click here to submit an issue on github.