【PAT甲级 - C++题解】1057 Stack

✍个人博客：https://blog.youkuaiyun.com/Newin2020?spm=1011.2415.3001.5343
📚专栏地址：PAT题解集合
📝原题地址：题目详情 - 1057 Stack (pintia.cn)
🔑中文翻译：栈
📣专栏定位：为想考甲级PAT的小伙伴整理常考算法题解，祝大家都能取得满分！
❤️如果有收获的话，欢迎点赞👍收藏📁，您的支持就是我创作的最大动力💪

1057 Stack

Stack is one of the most fundamental data structures, which is based on the principle of Last In First Out (LIFO). The basic operations include Push (inserting an element onto the top position) and Pop (deleting the top element). Now you are supposed to implement a stack with an extra operation: PeekMedian – return the median value of all the elements in the stack. With N elements, the median value is defined to be the (N/2)-th smallest element if N is even, or ((N+1)/2)-th if N is odd.

Input Specification:

Each input file contains one test case. For each case, the first line contains a positive integer N (≤10⁵). Then N lines follow, each contains a command in one of the following 3 formats:

Push key
Pop
PeekMedian

where key is a positive integer no more than 10⁵.

Output Specification:

For each Push command, insert key into the stack and output nothing. For each Pop or PeekMedian command, print in a line the corresponding returned value. If the command is invalid, print Invalid instead.

Sample Input:

17
Pop
PeekMedian
Push 3
PeekMedian
Push 2
PeekMedian
Push 1
PeekMedian
Pop
Pop
Push 5
Push 4
PeekMedian
Pop
Pop
Pop
Pop

Sample Output:

Invalid
Invalid
3
2
2
1
2
4
4
5
3
Invalid

题意

这题需要实现栈的基本操作，但这里除了栈的 Push 和 Pop 操作外，还需要我们额外实现一个操作，即取栈中所有元素的中位数，如果元素个数为偶数，则取第 N/2 个；如果元素个数为奇数，则取第 N/2+1 个。

其中 Push 操作不需要输出任何元素，而 Pop 操作需要输出删除的那个元素，取中位数操作则需要输出对应的中位数。另外，如果栈中没有元素，则 Pop 操作和取中位数操作输出 Invalid 。

思路

这道题其实是取数据流中位数那道题的扩展，没接触过该类型的小伙伴可以先去看看那道题目，传送门如下：

(369条消息) 剑指offer 42. 数据流中的中位数_Pandaconda的博客-优快云博客

而该题扩展到需要我们在取中位数的同时完成栈的两个基本操作，所以不能按照上面那道题中用两个对顶堆即大根堆和小根堆进行实现，因为 STL 提供的两个堆容器不能完成删除操作。

但我们可以模拟上面那题的做法，用两个 set 容器代替两个堆，但是题目可能会出现重复元素，所以我们这题需要用到两个 multiset 容器即假设为 up 和 down 。

上面的小根堆（存栈中更大的那一半数据）用 up 替代后，其堆顶元素可以用 up.begin() 来获取；而下面的大根堆（存栈中更小的那一半数据）用 down 替代后，其堆顶元素可以用 auto it=down.end(),it-- 来获取，因为没有相关函数可以直接获取 multiset 的队头和队尾元素，所以需要我们用迭代器进行模拟。

然后，我们需要保证 down 的元素个数始终要大于等于 up 的元素个数，这样取中位数的时候我们只用去取 down 中最大那个数即容器尾部那个数即可。

另外，需要注意的是还需要一个栈 stk 来存储元素，因为这两个容器只是方便我们取中位数，栈的插入和删除操作还需要在栈中执行。

所以，插入和删除操作后还要更新两个容器。接下来，我们就可以对三个操作进行模拟：

1. Push 操作，如果 up 为空的话或插入的元素 x 要小于 up 中的所有元素即 x<*up.begin() ，就将 x 插入 down 中；反之，插入 up 中。

2. Pop 操作，如果栈为空，则操作非法，需要输出 Invaild ；如果不为空，则需要判断删除的元素 x 在哪个容器中，如果 x 不在 down 中，则一定在 up 中（废话），所以执行对应容器的 erase 函数即可。

   但可能存在多个和 x 值相同的元素，如果直接 erase(x) 会将所有与 x 值相同的元素都删除，所以要用 find 函数返回指向 x 的迭代器，再执行 erase 函数后就只会删除迭代器指向的那一个元素了。

3. PeekMedian 操作，如果栈为空，则操作非法，需要输出 Invaild ；如果不为空，直接输出 down 容器的尾元素即可。

注意，每次 Push 和 Pop 操作后，可能会导致 up 和 down 中元素的个数不满足要求，所以每次执行完插入和删除操作后都需要调用一次 adjust 调整函数。

adjust 函数实际就执行两个操作：

1. 如果 down 的元素个数少于 up 的元素个数，就将 up 的队头元素移到 down 中。
2. 如果 down 的元素个数大于 up 的元素个数 + 1 ，就将 down 的队尾元素移到 up 中。

代码

#include<bits/stdc++.h>
using namespace std;

int n;
stack<int> stk;
multiset<int> up, down;

//down的元素个数要始终大于等于up的元素个数
void adjust()
{
    //down的元素个数不能少于up的元素个数
    while (up.size() > down.size())
    {
        auto it = up.begin();
        down.insert(*it);
        up.erase(it);
    }
    //down的元素个数不能大于up的元素个数+1
    while (down.size() > up.size() + 1)
    {
        auto it = down.end();
        it--;
        up.insert(*it);
        down.erase(it);
    }
}

int main()
{
    scanf("%d", &n);
    char op[20];
    for (int i = 0; i < n; i++)
    {
        scanf("%s", op);
        if (strcmp(op, "Push") == 0)    //插入操作
        {
            int x;
            scanf("%d", &x);
            stk.push(x);
            //down的元素要始终小于等于up的元素
            if (up.empty() || x < *up.begin())    down.insert(x);
            else up.insert(x);
            adjust();
        }
        else if (strcmp(op, "Pop") == 0)    //删除操作
        {
            if (stk.empty()) puts("Invalid");
            else
            {
                int x = stk.top();
                stk.pop();
                printf("%d\n", x);
                auto it = down.end();
                it--;
                //
                if (x <= *it)  down.erase(down.find(x));
                else up.erase(up.find(x));
                adjust();
            }
        }
        else    //输出中位数
        {
            if (stk.empty()) puts("Invalid");
            else
            {
                auto it = down.end();
                it--;
                printf("%d\n", *it);
            }
        }
    }

    return 0;
}