LC-895.最大频率栈

题目描述

leetcode 困难题

设计一个类似堆栈的数据结构，将元素推入堆栈，并从堆栈中弹出出现频率最高的元素。

实现 FreqStack 类:

• FreqStack() 构造一个空的堆栈。
• void push(int val) 将一个整数 val 压入栈顶。
• int pop() 删除并返回堆栈中出现频率最高的元素。
  • 如果出现频率最高的元素不只一个，则移除并返回最接近栈顶的元素。

示例1：

输入：
["FreqStack","push","push","push","push","push","push","pop","pop","pop","pop"],
[[],[5],[7],[5],[7],[4],[5],[],[],[],[]]
输出：[null,null,null,null,null,null,null,5,7,5,4]

提示1：

0 <= val <= 109
push 和 pop 的操作数不大于 2 * 104。
输入保证在调用 pop 之前堆栈中至少有一个元素。

哈希表 + 大根堆

比较复杂的做法是使用哈希表 $indexListMap$ 维护元素和出现过的下标集合之间的映射关系，再使用大根堆维护元素之间的排序关系：出现频率更高的在队列首部，频率相同时取下标更大的。

缺点是由于相同元素在优先队列中只存储一次，所以每次 $push$ 或者 $pop$ 操作的时候都需要将当前元素从有限队列中移除后再重新加入，因为此时优先关系可能会发生变化。

class FreqStack {
    private int index = 1;
    private final Map<Integer, List<Integer>> indexListMap = new HashMap<>();
    private final PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> {
        List<Integer> index1 = indexListMap.get(a);
        List<Integer> index2 = indexListMap.get(b);
        if(index1.size() != index2.size()){
            return index2.size() - index1.size();
        }
        return index2.get(index2.size() - 1) - index1.get(index1.size() - 1);
    });

    public void push(int val) {
        if(indexListMap.containsKey(val)){
            maxHeap.remove(val);
            indexListMap.get(val).add(index++);
            maxHeap.offer(val);
            return;
        }  
        List<Integer> indexList = new ArrayList(List.of(index++));
        indexList.add(index++);
        indexListMap.put(val, indexList);
        maxHeap.offer(val);
    }
    
    public int pop() {
        int val = maxHeap.poll();
        List<Integer> index = indexListMap.get(val);
        if(index.size() == 1){
            indexListMap.remove(val);
            return val;
        }
        index.remove(index.size() - 1);
        maxHeap.offer(val);
        return val;
    }
}

哈希表 + 大根堆 优化

更简单的做法是哈希表只维护元素和出现频率之间的关系，使用大根堆存储一个三元组 ${val, index, cnt}$ 分别代表元素本身、当前下标、当前频率，排序规则同上。

由于优先队列保存了完整的数据，即每个元素出现了几次就会出现在队列中几次，所以效率更高代码也更简洁。

class FreqStack {
    private int index = 1;
    private final Map<Integer, Integer> COUNT_MAP = new HashMap<>();
    // int{val, index, cnt}
    private final PriorityQueue<int[]> MAX = new PriorityQueue<>((a, b) -> b[2] != a[2] ? b[2] - a[2] : b[1] - a[1]);
    
    public void push(int val) {
        int count = COUNT_MAP.compute(val, (key, old) -> old == null ? 1 : ++old);
        MAX.offer(new int[]{val, index++, count});
    }
    
    public int pop() {
        int[] a = MAX.poll();
        COUNT_MAP.compute(a[0], (key, count) -> --count);
        return a[0];
    }
}

哈希表 + 栈

在上面的做法中，由于我们使用到了优先队列，所以每次入队和出队时间复杂度都为 $log(N)$ 。

另一种做法同样是使用哈希表维护元素和出现频率之间的关系，但是不再使用优先队列而是使用另一个哈希表来维护每种出现频率和对应的栈之间的关系，即会有多个栈。

稍微要注意的是，假设某个元素出现了 $5$ 次，那么在频率为 $1,2,3,4,5$ 的栈中都会有该元素，所以每当最大的栈为空时，我们只需要将最大栈的变量减一即可。

class FreqStack {
    private int max = 0;
    private final Map<Integer, Integer> freq = new HashMap<>();
    private final Map<Integer, Deque<Integer>> groupStack = new HashMap<>(); 
    
    public void push(int val) {
        int x = freq.compute(val, (key, old) -> old == null ? 1 : ++old);
        groupStack.compute(x, (key, old) -> {
            if(old == null){
                old = new ArrayDeque<>();
            }
            old.push(val);
            return old;
        });
        max = Math.max(max, x);
    }
    
    public int pop() {
        Deque<Integer> stack = groupStack.get(max);
        int val = stack.pop();
        if(stack.isEmpty()){
            max--; // 如果某个元素出现了 5 次，那么在频率为 1,2,3,4,5 的栈中都会有该元素
        }
        freq.compute(val, (key, old) -> --old);
        return val;
    }
}

哈希表 + List

类似上一种做法，可以使用哈希表加上 $List$ 来实现。

class FreqStack {
    private final Map<Integer, Integer> freq = new HashMap<>();
    private final List<Deque<Integer>> stackList = new ArrayList();
    
    public void push(int val) {
        int x = freq.getOrDefault(val, 0);
        if(x == stackList.size()){
            stackList.add(new ArrayDeque<>());
        }
        stackList.get(x).push(val);
        freq.put(val, x + 1);
    }
    
    public int pop() {
        Deque<Integer> last = stackList.get(stackList.size() - 1);
        int val = last.pop();
        if(last.isEmpty()){
            stackList.remove(stackList.size() - 1);
        }
        freq.compute(val, (__, old) -> --old);
        return val;
    }
}