【PAT】A1057 Stack (30分)（分块思想 / 树状数组）

1057 Stack (30分)

题目链接

思路：分块的思想，详细略。

#include<iostream>
#include<cmath>
#include<algorithm>
#include<stack>
#include<cstring>
using namespace std;

const int maxn = 100005;

int table[maxn] = {0};
int *block;
stack<int >st;
int blockCounts, blockLen;//blockCounts:10^5分的块数， blockLen:每块的块长 

int findMedian()
{
	int stLen = (st.size() + 1)/2;
//	if(stLen & 1)//奇数 
//		stLen = (stLen+1)/2;
//	else
//		stLen /= 2;
	int sum = 0;
	for(int i=0;i<blockLen;i++){
		if(sum + block[i] < stLen)//不在第i块内 
			sum += block[i];
//		else if(sum + block[i] == stLen){//在这一个块中的最后一个 
//			for(int j = (i+1)*blockLen -1;;j--)
//				if(table[j]!=0)
//					return j; 
//		}
		else{//sum + block[i] >= stLen
			//在block[i]中 
			for(int j=i*blockLen;j<(i+1)*blockLen;j++)
				if(sum + table[j] < stLen)
					sum += table[j];
				else
					return j;
		}
	}
}

int main()
{
	std::ios::sync_with_stdio(false);
	int n;
	cin>>n;
	blockCounts = (int)sqrt(1.0 * maxn) + 1;//分块数 
	blockLen = (int)sqrt(1.0 * maxn);//每块的块长 
	block = new int[blockCounts];
	memset(block, 0, sizeof(int)*blockCounts);
	string temp;
	int key;
	for(int i=0;i<n;i++)
	{
		cin>>temp;
		if(temp == "Pop")
		{
			if(st.size() == 0)
				cout<<"Invalid"<<endl;
			else{
				cout<<st.top()<<endl;
				table[st.top()]--;
				block[st.top() / blockLen]--;
				st.pop();
			}
		}
		else if(temp == "PeekMedian")
		{
			if(st.size() == 0)
				cout<<"Invalid"<<endl;
			else
				cout<<findMedian()<<endl;
		}
		else{//push
			cin>>key;
			table[key]++;
			block[key / blockLen]++;
			st.push(key);
		}
	}
	return 0;
}

本来想用快排来做，但是想了想不行，肯定会超时，虽然快排的复杂度为O(nlogn)，但是如果使用数组来存的话，删除元素复杂度是O(n)，肯定会超时。

---

更新：对上述代码做了一些优化。

#include<iostream>
#include<cmath>
#include<algorithm>
#include<stack>
#include<cstring>
using namespace std;

const int maxn = 100005;

int table[maxn] = {0};
int *block;
stack<int >st;
int blockCounts, blockLen;//blockCounts:10^5分的块数， blockLen:每块的块长 

int findMedian(int k)//查找第K大的数 
{
	int sum = 0, idx;
	while(sum + block[idx] < k)
		sum += block[idx++];
	int firstIdx = idx * blockLen;
	while(sum + table[firstIdx] < k)
		sum += table[firstIdx++];
	return firstIdx;
}

int main()
{
	std::ios::sync_with_stdio(false);
	int n;
	cin>>n;
	blockCounts = (int)sqrt(1.0 * maxn) + 1;//分块数（向上取整） 
	blockLen = (int)sqrt(1.0 * maxn);//每块的块长（向下取整） 
	block = new int[blockCounts];
	memset(block, 0, sizeof(int)*blockCounts);
	string temp;
	int key;
	for(int i=0;i<n;i++)
	{
		cin>>temp;
		if(temp == "Pop")
		{
			if(st.size() == 0)
				cout<<"Invalid"<<endl;
			else{
				cout<<st.top()<<endl;
				table[st.top()]--;
				block[st.top() / blockLen]--;
				st.pop();
			}
		}
		else if(temp == "PeekMedian")
		{
			if(st.size() == 0)
				cout<<"Invalid"<<endl;
			else
				cout<<findMedian((st.size() + 1)/2)<<endl;
		}
		else{//push
			cin>>key;
			table[key]++;
			block[key / blockLen]++;
			st.push(key);
		}
	}
	return 0;
}

---

再次更新：使用树状数组的思想来做，也即单点更新、区间查询的思想。

#include<iostream>
#include<cmath>
#include<algorithm>
#include<stack>
#include<cstring>
using namespace std;

const int maxn = 100005;
#define lowbit(i) (i & (-i))

int c[maxn];//全局变量会自动初始化为全0
int maxNum = -1;//记录目前最大的数 ,可能会有点用 
stack<int >st;

//update函数将第x个整数上加上v
void update(int x, int v)
{
	//i必须是<maxn,不能是maxNum
	for(int i=x; i<maxn; i+=lowbit(i))
		c[i] += v;
}

//getSum函数返回前x个整数之和
int getSum(int x)
{
	int sum = 0;
	for(int i=x; i>0; i-=lowbit(i))
		sum += c[i];
	return sum;	
}

//二分法查找第k大的元素 
int findKthElement(int k)
{
	int low = 1, high = maxNum;//high = maxn也是可以的
	while(low < high)
	{
		int mid = (low + high) / 2;
		if(getSum(mid) >= k)//一定在mid之内 即<=mid（包括mid） 
			high = mid;
		else//getSum(mid) < k //一定大于mid 
			low = mid+1; 
	}
	return low;
}

//顺序遍历查找第k大的元素 
//超时 
int findMedian(int k)
{
	int low = 0;
	while(1)
	{
		if(getSum(low) < k)
			low++;
		else
			return low;
	}
}

int main()
{
	std::ios::sync_with_stdio(false);
	int n;
	cin>>n;
	string temp;
	int key;
	for(int i=0;i<n;i++)
	{
		cin>>temp;
		if(temp == "Pop")
		{
			if(st.size() == 0)
				cout<<"Invalid"<<endl;
			else{
				cout<<st.top()<<endl;
				update(st.top(), -1);//树状数组c[st.top()]更新-1
				st.pop();
			}
		}
		else if(temp == "PeekMedian")
		{
			if(st.size() == 0)
				cout<<"Invalid"<<endl;
			else
//				cout<<findMedian((st.size() + 1)/2)<<endl; //顺序遍历查找第k大的数 ----超时 
				cout<<findKthElement((st.size() + 1)/2)<<endl; //二分查找第k大的数 
		}
		else{//push
			cin>>key;
			if(key > maxNum)
				maxNum = key;
			update(key, 1);//树状数组c[key]更新1
			st.push(key);
		}
	}
	return 0;
}

注意：在查找的时候，使用顺序遍历的方法是会超时的。使用二分的方法就不会。