题目
Design a max stack that supports push, pop, top, peekMax and popMax.
push(x) -- Push element x onto stack.
pop() -- Remove the element on top of the stack and return it.
top() -- Get the element on the top.
peekMax() -- Retrieve the maximum element in the stack.
popMax() -- Retrieve the maximum element in the stack, and remove it. If you find more than one maximum elements, only remove the top-most one.
不难的题目,但是写起来巨烦,因为java的优先队列要自己写Compare函数。。
思路很简单
用一个stack存数据,再用一个优先队列存一个pair <value, index in stack>, 加上lazy removal。
答案
/*
Use a normal stack
Then, use a heap to keep track of the maximum guy, and its index
*/
class MaxStack {
class Entry implements Comparable<Entry> {
int val;
int idx;
public Entry(int val, int idx) {
this.val = val;
this.idx = idx;
}
@Override
public int compareTo(Entry o) {
if(this.val < o.val) return 1;
else if(this.val > o.val) return -1;
else {
if(this.idx < o.idx) return 1;
else if(this.idx > o.idx) return -1;
else return 0;
}
}
@Override
public boolean equals(Object obj) {
Entry e = (Entry) obj;
return this.val == e.val && this.idx == ((Entry) obj).idx;
}
}
Stack<Integer> stk;
PriorityQueue<Entry> pq;
Map<Integer, Integer> map;
Set<Integer> set;
/** initialize your data structure here. */
public MaxStack() {
stk = new Stack<Integer>();
pq = new PriorityQueue<Entry>();
set = new HashSet<>();
}
public void push(int x) {
stk.push(x);
pq.offer(new Entry(x, stk.size() - 1));
//pq.offer(new int[]{x, stk.size() - 1});
}
public int pop() {
// Look at current index
int idx = stk.size() - 1;
while(set.contains(idx)) {
stk.pop();
set.remove(idx);
idx = stk.size() - 1;
}
pq.remove(new Entry(stk.peek(), stk.size() - 1));
return stk.pop();
}
public int top() {
// Look at current index
int idx = stk.size() - 1;
while(set.contains(idx)) {
stk.pop();
set.remove(idx);
idx = stk.size() - 1;
}
return stk.peek();
}
public int peekMax() {
return pq.peek().val;
}
public int popMax() {
// Add this to map, lazy delete
Entry e = pq.poll();
set.add(e.idx);
return e.val;
}
}