(PAT Advanced)1057 Stack（树状数组+二分查找） C++

原题：https://pintia.cn/problem-sets/994805342720868352/problems/994805417945710592

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

题目大意

实现一种特殊的堆栈结构，除了满足常规栈结构的操作外，还需要能够返回栈中元素的中指。

暴力解法（Hash映射 超时

使用一个数组book[]来存放每个数字出现的次数，当需要找中值的时候，就全局遍历一次数组book[]，找到位于mid位置的元素，但是由于总的元素个数可能在$10^5$左右，所以有两个测试点超时。

/* 1057 Stack (30 分) */
// 暴力方法，使用hash映射输入模拟
#include<iostream>
#include<string.h>
#include<stack>

using namespace std;

const int maxlen = 100010; // 最多的栈的元素
int book[maxlen];  // 存储每个元素出现的次数
int main() {
	fill(book, book + maxlen, 0);
	int t, tmp;
	scanf("%d", &t);
	char op[15];
	stack<int> S;
	for (int i = 0; i < t; i++) {
		scanf("%s", op);
		if (op[1] == 'o') {
			if (!S.size()) printf("Invalid\n");
			else {
				printf("%d\n", S.top());
				book[S.top()]--; S.pop();
			}
		}
		else if (op[1] == 'u') {
			scanf("%d", &tmp);  book[tmp]++;
			S.push(tmp);
		}
		else {
			if (!S.size()) {
				printf("Invalid\n"); continue;
			}
			int mid = (S.size() + 1) / 2;
			for (int i = 1; i <= maxlen; i++) {
				if (book[i] > 0) mid -= book[i];
				if (mid <= 0) {
					printf("%d\n", i); // 包含第mid个元素
					break;
				}
			}
		}
	}
	system("pause");
	return 0;
}

优化：树状数组+二分查找

树状数组适用于单点更新和区间查询的情况，且插入和查询的复杂度都在$O(logn)$，所以我们假设数组a[]来记录每个数字出现的次数，树状数组book[i]就表示a[1]~a[i]出现的累计次数，然后我们在1~maxn区间内进行二分查找来找恰好累计为中值的数字。

/* 1057 Stack (30 分) */
// 使用树状数组来存储数字出现的次数，用二分查找来查找第k大的元素
#include<iostream>
#include<string.h>
#include<stack>

using namespace std;

#define lowbit(i) (i)&(-i)
const int maxlen = 100010; // 最多的栈的元素

stack<int> S;
int book[maxlen];  // 存储每个元素出现的次数，没出现一个标记以及
void update(int x, int v) {
	// 表明数字x出现
	for (int i = x; i < maxlen; i += lowbit(i)) book[i] += v;
}

int getsum(int x) {
	int sum = 0;
	for (int i = x; i; i -= lowbit(i)) sum += book[i];
	return sum;
}

void peekMedian() {
	// 求第k大的数字
	int left = 1, right = maxlen, mid, k = (S.size() + 1) / 2;
	while (left < right) {
		mid = (left + right) / 2;
		if (getsum(mid) >= k)
			right = mid; // 可以恰好等于或者包含第k个元素
		else
			left = mid + 1;
	}
	printf("%d\n", left);
}
int main() {
	fill(book, book + maxlen, 0);
	int t, tmp;
	scanf("%d", &t);
	char op[15];
	for (int i = 0; i < t; i++) {
		scanf("%s", op);
		if (op[1] == 'o') {
			if (!S.size()) printf("Invalid\n");
			else {
				update(S.top(), -1);
				printf("%d\n", S.top());
				S.pop();
			}
		}
		else if (op[1] == 'u') {
			scanf("%d", &tmp);  update(tmp,1);
			S.push(tmp);
		}
		else {
			if (!S.size())
				printf("Invalid\n");
			else peekMedian();
		}
	}
	system("pause");
	return 0;
}