PAT.1057 Stack

PAT.1057 Stack

题目链接
看到题，是个提供能求中位数的栈，好哇直接模拟，然后超时。
一开始想到的方法使用multiset来维护栈内数据，然后通过迭代器来获取中位数，但这样实质上效率很差。
于是看了一些别人的解法，发现树状数组可以做，但是呢，我不会树状数组，付出了代价。
后来看到了用multiset维护两段数据的写法，链接。
于是按照这个思路写出了题解：用upper维护大于中位数的元素，用lower维护小于等于中位数的元素，两者都是multiset，begin()指向贴近中位数的数，每次栈中元素发生变化的时候对两侧元素进行修正。
修正思路如下：

• 123 455 - 1233 455 不用调整 - 123 4455 要调整
• 123 45 - 1233 45 要调整 - 123 455不用调整

因此可以得到，当元素变动后upper元素个数大于lower时需要调整，当lower元素个数大于upper元素个数+1时需要调整。

错误的

#include<bits/stdc++.h>

using namespace std;
typedef long long ll;

int n,num;
string cmd;
stack<int> stk;
multiset<int> nums;

int main(){
    cin>>n;
    for(int i = 0 ; i < n ; ++i){
        cin>>cmd;
        if(cmd == "Push"){
            cin>>num;
            stk.push(num);
            nums.insert(num);
        }else if(cmd == "Pop"){
            if(stk.size() < 1) printf("Invalid\n");
            else{
                printf("%d\n",stk.top());
                nums.erase(nums.find(stk.top()));
                stk.pop();
            }
        }else if(cmd == "PeekMedian"){
            if(stk.size() < 1) printf("Invalid\n");
            else{
                int diff = (stk.size() + 1) / 2 - 1;
                auto pos = nums.begin();
                for(int i = 0 ; i < diff ; ++i) pos++;
                cout<<*pos<<endl;
            }
        }
    }
}

正确的

#include<bits/stdc++.h>

using namespace std;
typedef long long ll;

int n,num,median;
string cmd;
stack<int> stk;
multiset<int> upper;  
multiset<int,greater<int>> lower;  

void getMedian(){
    if(upper.size() > lower.size()){
        lower.insert(*upper.begin());
        upper.erase(upper.begin());
    }else if(lower.size() > upper.size() + 1){
        upper.insert(*lower.begin());
        lower.erase(lower.begin());
    }
    median = *(lower.begin());
}

int main(){
    cin>>n;
    for(int i = 0 ; i < n ; ++i){
        cin>>cmd;
        if(cmd == "Push"){
            cin>>num;
            stk.push(num);
            if(num > median) upper.insert(num);
            else lower.insert(num);
            getMedian();
        }else if(cmd == "Pop"){
            if(stk.size() < 1) printf("Invalid\n");
            else{
                printf("%d\n",stk.top());
                if(stk.top() > median) upper.erase(upper.find(stk.top()));
                else lower.erase(lower.find(stk.top()));
                stk.pop();
                getMedian();
            }
        }else if(cmd == "PeekMedian"){
            if(stk.size() < 1) printf("Invalid\n");
            else printf("%d\n",median);
        }
    }
}