PAT(A) - 1057. Stack (30)

1057. Stack (30)

时间限制

150 ms

内存限制

65536 kB

代码长度限制

16000 B

判题程序

Standard

作者

CHEN, Yue

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操作，额外增加一个PeekMedian操作，意为返回栈中第K小的元素，K = N / 2（当N是偶数），K = （N + 1）/ 2（当K是奇数）。这里N为当前栈中元素的个数。输入就是一系列操作，其中Push操作后面要加上一个数字，代表压栈的元素。输出为一系列操作的结果，如果操作为Push什么也不输出；如果是Pop，输出弹出元素的值；如果是PeekMedian，输出相应的值；如果操作非法（比如栈空还执行Pop或PeekMedian），输出“Invalid”。

我刚开始的做法是模拟，用STL中的stack进行压栈弹栈，vector也有push_back和pop_back操作，用来存储当前栈中的所有元素。另外又开辟了一个很大的数组a，每次执行PeekMedian操作，vector把所有元素赋值给a，a再进行排序，找到第K小的元素。这种做法能达到目的，但是显然很暴力。这种方案只拿到了17分。而且这道题时间卡的很严格，150ms，显然我们的方案不可取。

后来实在没招了，想过用队列、或再开一个栈或者是建一棵二叉搜索树，但想了想感觉都达不到题目要求。。只好看别人做的了，大家几乎都是用树状数组，以前听过这个词，但没了解是怎么用的。不清楚树状数组的可以网上搜一下哈~

#include <cstdio>
#include <stack>
#include <string.h>
#include <vector>
#include <algorithm>

using namespace std;

const int Max_required = 100005;
int tree_array[Max_required];

inline int lowbit(int x)
{
    return x&(-x);
}

void add(int x, int value)  //更新前n项和函数
{
    while (x < Max_required)
    {
        tree_array[x] += value;
        x += lowbit(x);
    }
}

int sum(int x)      //求前n项和函数
{
    int total = 0;
    while (x > 0)
    {
        total += tree_array[x];
        x -= lowbit(x);
    }
    return total;
}

int binary_find(int x)  //二分查找代码
{
    int low = 1, high = Max_required, mid;

    while (low <= high)
    {
        mid = (low + high) >> 1;
        int total = sum(mid);

        if (total >= x)
            high = mid - 1;
        else if (total < x)
            low = mid + 1;
    }
    return low;
}

int main() {

    stack<int> s;

    int n;
    scanf( "%d", &n );
    for( int i = 0; i < n; i++ ) {
        char str[15];
        int key;
        scanf( "%s", str );
        if( !strcmp( str, "Push" ) ) {
            scanf( "%d", &key );
            s.push( key );
            add( key, 1 );
            continue;
        }
        if( !strcmp( str, "Pop" ) ) {
            if( s.empty() ) {
                printf( "Invalid\n" );
            }
            else {
                int top = s.top();
                s.pop();
                add( top, -1 );
                printf( "%d\n", top );
            }
            continue;
        }
        if( !strcmp( str, "PeekMedian" ) ) {
            if( s.empty() ) {
                printf( "Invalid\n" );
            }
            else {
                int len = s.size();
                int res = -1;
                if ( len % 2 == 1 )
                    res = binary_find( ( len + 1 ) / 2 );
                else
                    res = binary_find( len / 2 );
                printf( "%d\n", res );
            }
        }
    }
    return 0;
}