1057. Stack (30) - 树状数组

Posted 2020-09-23 yfceshi

题目例如以下：

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

这个题目我最初用的是string、stringstream和vector来做。发现会严重超时，后来在网上參考了和山米兰的解法。发现他的方法非常有技巧，分析例如以下：

①对命令的解析，仅仅看第二位，假设是o，说明是Pop，假设是e。说明是PeekMedian，否则是push，是push则应当再读入一次数字，这比用getline要好的多，由于getline还须要排除第一个输入的N。

②求中位数的思想，不是排序找中间的值，而是通过统计从1開始的每一个元素的个数放到数组C中，这样从前到后。数组C的子列和为题目要求的位置时，拿到的就是中位数。

③求子列和的思想。由于是从前到后的前缀和，能够利用树状数组。以下的代码利用add实现了加入和删除两种操作。利用value的不同，1表示加入。2表示删除。树状数组的基本思想就是数组C中不同元素管辖不同的区域。假设要加入一个元素。则全部满足区域条件的位置都要+value。反之假设删除。全部满足条件的区域都要-value。本题要求的是统计1~100000的元素个数。因此value=+1或者-1。

④求子列和为题目要求的值，利用二分查找。

#include<stdio.h>
#include<cstring>
#include<iostream>
#include<string>
using namespace std;

const int N=100001;
int c[N];

int lowbit(int i){
	return i&(-i);
}

void add(int pos,int value){
	while(pos<N){
		c[pos]+=value;
		pos+=lowbit(pos);
	}
}

int sum(int pos){
	int res=0;
	while(pos>0){
		res+=c[pos];
		pos-=lowbit(pos);
	}
	return res;
}

int find(int value){
	int l=0,r=N-1,median,res;
	while(l<r-1){
		if((l+r)%2==0)
			median=(l+r)/2;
		else
			median=(l+r-1)/2;
		res=sum(median);
		if(res<value)
			l = median;
		else
			r = median;

	}
	return l+1;
}

int main(){
	char ss[20];
	int stack[N],top=0,n,pos;
	memset(c,0,sizeof(c));
	scanf("%d",&n);
	while(n--){
		scanf("%s",ss);
		if(ss[1]==‘u‘){
			scanf("%d",&pos);
			stack[++top]=pos;
			add(pos,1);
		}else if(ss[1]==‘o‘){
			if(top==0){
				printf("Invalid\n");
				continue;
			}
			int out=stack[top];
			add(out,-1); // 删除元素out
			printf("%d\n",stack[top--]);
		}else if(ss[1]==‘e‘){
			if(top==0){
				printf("Invalid\n");
				continue;
			}
			int res;
			if(top%2==0)
				res=find(top/2);
			else
				res=find((top+1)/2);
			printf("%d\n",res);

		}else{
			printf("Invalid\n");
		}
	}

	return 0;
}