πŸ“ Problem Details

πŸ’­What Were My Initial Thoughts?

- private class members for:
	- capacity
	- vector of stacks to allow for infinite addition of stacks
	- need a way of tracking available positions in the event of popAtStack

πŸ€”What Did I Struggle With?

- coming up with the idea to use a set to track available positions 

πŸ’‘ Explanation of Solution

capacity tracking:
	- store the capacity as a private class member so all methods can access it 

vector of stacks:
	- use a vector<stack<int>> to allow for dynamically expanding stacks 

track available stacks:
	- to optimize the push function, track the first available stack that can take a new value 
	- this can be achieved using a set<int> of indicies for non-full stacks
	- this ensures push remains efficient 

track active stacks:
	- for the pop function, track non-empty stacks 

βŒ› Complexity Analysis

O(log N) for push
O(log N) for popAtStack 

πŸ’» Implementation of Solution

class DinnerPlates {
 
private:
    int capacity;
    vector<stack<int>> stacks;
    set<int> available; // indicies of stacks that can accept more plates
 
public:
    DinnerPlates(int capacity) : capacity(capacity) {}
 
    void push(int val) {
 
        while (!available.empty() && (*available.begin() >= stacks.size() || stacks[*available.begin()].size() == capacity)) {
            available.erase(available.begin()); // Remove invalid or full indices
        }
 
        if (available.empty()) {
            stacks.push_back(stack<int>());
            available.insert(stacks.size() - 1);
        }
 
        int index = *available.begin();
        // Validate index
        if (index < 0 || index >= stacks.size()) {
            throw runtime_error("Invalid index detected in available set.");
        }
 
        stacks[index].push(val);
        
        if (stacks[index].size() == capacity) {
            available.erase(index);
        }
    }
 
    int pop() {
        while(!stacks.empty() && stacks.back().empty()) {
            stacks.pop_back();
        }
 
        if (stacks.empty()) return -1; // all stacks are empty
 
        int val = stacks.back().top();
 
        stacks.back().pop();
 
        if (stacks.back().size() < capacity) {
            available.insert(stacks.size() - 1);
        }
        return val;
    }
 
    int popAtStack(int index) {
        if (index < 0 || index >= stacks.size() || stacks[index].empty()) {
            return -1;
        }
        int val = stacks[index].top();
        stacks[index].pop();
 
        if (stacks[index].size() < capacity) {
            available.insert(index);
        }
        return val;
    }
};
 
/**
 * Your DinnerPlates object will be instantiated and called as such:
 * DinnerPlates* obj = new DinnerPlates(capacity);
 * obj->push(val);
 * int param_2 = obj->pop();
 * int param_3 = obj->popAtStack(index);
 */