LeetCode 895. Maximum Frequency Stack
题目链接:https://leetcode.com/problems/maximum-frequency-stack/
题意:实现一种数据结构FreqStack,FreqStack需要实现两个功能:
- push(int x) : 将x放入栈中
- pop(): 移除并返回栈中出现次数最多的元素,若果该元素不唯一,返回离栈顶最近的
思路:维护一个set,set中存储元素值,以及一个包含该元素在原栈中的出现次序的栈,排序先按照set中栈的大小排序,栈大小相同则按照栈顶元素大小排序。push若该元素不在set中,则直接插入,否则找到该元素,往栈中加入一个新次序。pop时只需弹出set的起始位置栈中的第一个元素。为了便于push时找到该元素,使用unordered_map记录元素在set中的位置。时间复杂度$O(log n)$
class FreqStack { public: struct P { int data; //存储入栈的元素,相同的元素在一起 stack< int > q; //存储元素的插入顺序 bool operator < (const P &a) const { if (q.size() != a.q.size()) //先按元素个数排序 return q.size() > a.q.size(); return q.top() > a.q.top(); //再按离栈顶顺序排序 } P(int x) { data = x; } }; set<P> s; int cnt = 0; unordered_map<int, set<P>::iterator > mp; //存储元素在set中的位置 void push(int x) { P a(x); cnt++; //cnt表示入栈时间 if (mp.find(x) != mp.end()) { //元素在set中已经存在 cnt++; a = *mp[x]; s.erase(mp[x]); } a.q.push(cnt); s.insert(a); mp[x] = s.find(a); } int pop() { P a = *s.begin(); //set中第一个 mp.erase(a.data); s.erase(s.begin()); a.q.pop(); if (a.q.size() > 0) { s.insert(a); mp[a.data] = s.find(a); } return a.data; } };
官方的思路比较巧妙,复杂度$O(1)$:
维护两个hash表time<int, int>和mp<int, stack<int> >以及一个max_time,time[i]记录元素i出现的次数,mp[i]记录出现次数>=i次的元素,按照出现次序入栈,max_time记录当前出现次数最多元素的次数。
push(i)时time[i]++且更新max_time,pop时将出现次数为max_time的栈中第一个元素出栈,并更新max_time
class FreqStack { public: void push(int x) { time[x]++; if(time[x]>max_time) max_time=time[x]; mp[time[x]].push(x); //将x出现次数加入对应的堆栈中 } int pop() { int z = mp[max_time].top(); //出现最多次数的栈顶元素 mp[max_time].pop(); if(mp[max_time].empty()) //更新max_time max_time--; time[z]--; return z; } private: unordered_map<int,int> time; //元素出现的次数 unordered_map<int,stack<int> > mp; //出现次数对应的元素 int max_time=0; };