classDinnerPlates { vector<stack<int>> stacks; set<int> condidates; int capacity; public: DinnerPlates(int capacity_) : capacity(capacity_) { } voidpush(int val){ if (condidates.empty()) { condidates.insert(stacks.size()); stacks.push_back(stack<int>()); } int target = *condidates.begin(); stacks[target].push(val); if (stacks[target].size() == capacity) { condidates.erase(target); } } intpop(){ if (stacks.empty()) return-1; int ret = stacks.back().top(); stacks.back().pop(); while (!stacks.empty() && stacks.back().empty()) { stacks.pop_back(); condidates.erase(stacks.size()); } if (!stacks.empty() && stacks.back().size() < capacity) { condidates.insert(stacks.size() - 1); } return ret; } intpopAtStack(int index){ if (index >= stacks.size()) return-1; if (stacks[index].empty()) return-1; if (index == stacks.size() - 1) returnpop(); int ret = stacks[index].top(); stacks[index].pop(); condidates.insert(index); return ret; } };
/** * 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); */