天梯赛L3-002 特殊堆栈 （树状数组+二分答案）

原题地址：https://pintia.cn/problem-sets/994805046380707840/problems/994805053695574016

L3-002 特殊堆栈 （30 分）

堆栈是一种经典的后进先出的线性结构，相关的操作主要有“入栈”（在堆栈顶插入一个元素）和“出栈”（将栈顶元素返回并从堆栈中删除）。本题要求你实现另一个附加的操作：“取中值”——即返回所有堆栈中元素键值的中值。给定 N 个元素，如果 N 是偶数，则中值定义为第 N/2 小元；若是奇数，则为第 (N+1)/2 小元。

输入格式：

输入的第一行是正整数 N（≤10​5​​）。随后 N 行，每行给出一句指令，为以下 3 种之一：

Push key
Pop
PeekMedian

其中 key 是不超过 10​5​​ 的正整数；Push 表示“入栈”；Pop 表示“出栈”；PeekMedian 表示“取中值”。

输出格式：

对每个 Push 操作，将 key 插入堆栈，无需输出；对每个 Pop 或 PeekMedian 操作，在一行中输出相应的返回值。若操作非法，则对应输出 Invalid。

输入样例：

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

输出样例：

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

解题思路

 常规栈的操作用stack实现就好了，主要讲一下取中值的操作。如果栈里有n个数，如果n是偶数，那中值就是n/2小元，就是第n/2小的数，如果n是奇数，就是第（n+1）/ 2小的数。因为key的范围是1到100000，所以，假设a[i]是用来记录i在栈中的个数，当将i入栈时，a[i]++，将i出栈时，a[i]--。a[]的前i个数的和就是小于等于i的数的个数。

用树状数组tree[]来维护a[]，query（i）返回a[]的前i项和。假设要找的是第z小的数，不一定query(i)等于z，i就是第z小。i可能不在栈里，即a[i]为0。是要query(i)大于等于z时的最小i，才是答案（大于等于是因为a[i]可能大于1）。所以要用二分答案来求解。

代码如下

#include <iostream>
#include <string>
#include <stack>
#include <cstring>
using namespace std;
int tree[100005];
int lowbit(int x)
{
	return x & (-x);
}
void update(int x, int k)   //更新 
{
	for(int i = x; i <= 100000; i += lowbit(i))
		tree[i] += k;
}

int query(int x)   //查询前x项和 
{
	int sum = 0;
	for(int i = x; i > 0; i -= lowbit(i))
		sum += tree[i];
	return sum;
}

int main()
{
	int n;
	while(cin >> n){
		memset(tree, 0, sizeof(tree));
		stack<int> sta;
		for(int i = 0; i < n; i ++){
			string str;
			cin >> str;
			if(str == "Push"){
				int key;
				cin >> key;
				sta.push(key);
				update(key, 1);
			}
			else if(str == "Pop"){
				if(!sta.empty()){
					cout << sta.top() << endl;
					update(sta.top(), -1);
					sta.pop();	
				}
				else
					cout << "Invalid" << endl;	
			}
			else {
				if(!sta.empty()){
					int n = sta.size();
					if(n % 2 == 0)
						n /= 2;
					else
						n = (n + 1) / 2;
					int l = 1;
					int r = 100000;
					int mid = (l + r) / 2;
					while(l != r){          //二分找第n小的数 
						if(query(mid) < n)
							l = mid + 1;
						else
							r = mid;
						mid = (l + r) / 2;
					}
					cout << mid << endl;
				}
				else 
					cout << "Invalid" << endl;
			}
		}
	}
	return 0;
}