Codechef SEAARC Sereja and Arcs (分块、组合计数)

我现在真的什么都不会了呢......

题目链接: https://www.codechef.com/problems/SEAARC

好吧，这题其实考察的是枚举的功力……

题目要求的是\(ABAB\)的数量，这个不太好求，但是不同颜色区间对的总数和\(AABB,ABBA\)的都比较好求

补集转化，求\(ans0,ans1,ans2\), 分别表示总数、\(AABB\)、\(ABBA\)的数量

\(ans0\)很好算

\(ans1\), 枚举\(B\)的左端点

\(ans2\), 分块讨论

若\(A,B\)都是小颜色(该颜色块的数量小于阈值)，则枚举\(A\)的左右端点，变成二维数点

若\(A\)是大颜色\(B\)是小颜色，则枚举\(A\)的种类，然后枚举\(B\)的种类和左右端点，推一下式子发现可以维护前缀和省去枚举左端点

若\(B\)是大颜色，同理枚举\(B\)的种类然后枚举\(A\)的种类和左右端点，同样省去枚举左端点

从昨天下午做到今天上午。。

理论最优时间复杂度\(O(n\sqrt{n\log n})\)

代码

#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<cassert>
#include<vector>
#include<cmath>
#define llong long long
using namespace std;

const int N = 1e5;
const int P = 1e9+7;
const llong INV2 = 5e8+4;
int nxt[N+3];
int lstpos[N+3];
int a[N+3];
int num[N+3];
int cnum[N+3];
llong tmp0[N+3],tmp1[N+3],tmp2[N+3];
vector<int> clrpos[N+3];
int n,m,B;
llong ans1,ans2a,ans2b,ans2c,ans0,ans;

llong C2(llong x) {return x*(x-1ll)/2ll%P;}
llong update(llong &x,llong y) {x = (x+y)%P;}

struct BITree
{
	llong tr[N+3]; int siz;
	void addval(int lrb,llong val)
	{
		while(lrb<=siz)
		{
			update(tr[lrb],val);
			lrb += (lrb&(-lrb));
		}
	}
	llong querysum(llong rb)
	{
		llong ret = 0ll;
		while(rb)
		{
			update(ret,tr[rb]);
			rb -= (rb&(-rb));
		}
		return ret;
	}
	void clear()
	{
		for(int i=0; i<=siz; i++) tr[i] = 0ll;
	}
} bit1,bit2;

void getans0() //ans0=ËùÓÐÑÕÉ«C(num,2)Á½Á½³Ë»ýÖ®ºÍ ²»ËãAAAA
{
	llong cur = 0ll;
	for(int i=1; i<=m; i++)
	{
		llong tmp = C2(num[i]);
		update(ans0,cur*tmp%P);
		update(cur,tmp);
	}
}

void getans1() //²»ËãAAAA
{
	llong tmp = 0ll; //µ±Ç°×ܹ²Í¬É«Çø¼äµÄ¸öÊý
	for(int i=1; i<=n; i++)
	{
		update(ans1,(tmp-C2(cnum[a[i]]))*(num[a[i]]-cnum[a[i]]-1)); //ÕâÖÖÑÕÉ«×ܸöÊý¼õÕâ֮ǰµÄ¸öÊý-1£¬µÈÓÚÕâÖ®ºóµÄ¸öÊý£¬³ËÒÔÇ°ÃæµÄÒìÉ«¸öÊý
		update(tmp,(llong)cnum[a[i]]); //ÕâÖÖÑÕɫ֮ǰµÄ¸öÊý
		cnum[a[i]]++;
	}
}

void getans2a() //small-small ²»ËãAAAA
{
	bit1.siz = n; bit1.clear(); llong cur = 0ll;
	for(int i=1; i<=n; i++) //ö¾ÙÓҶ˵ã
	{
		if(num[a[i]]<=B)
		{
			int tnum = 0;
			for(int j=nxt[i]; j; j=nxt[j])
			{
				llong tmp = cur-bit1.querysum(j)-C2(tnum)+P+P; //>jµÄ×ó¶ËµãµÄ¸öÊý=×ܸöÊý¼õ<=jµÄ¸öÊý£¬È¥µôͬɫµÄ¸öÊý(²»ËãiºÍj)
				update(ans2a,tmp);
				tnum++; //ͬɫ¸öÊý
			}
			for(int j=nxt[i]; j; j=nxt[j]) //ÓÉÓÚÊÇÓҶ˵ãСÓÚj£¬ÏÈÐ޸ĺó²éѯ
			{
				cur++; //Ä¿Ç°Çø¼ä×ܸöÊý,¼´bit1.querysum(i)
				bit1.addval(j,1);
			}
		}
	}
}

void getans2b() //large-small
{
	for(int i=1; i<=m; i++) //ö¾ÙAÖÖÀà
	{
		if(num[i]>B)
		{
			tmp1[0] = 0ll; for(int j=1; j<=n; j++) tmp1[j] = tmp1[j-1]+(a[j]==i?1:0);
			for(int j=1; j<=m; j++)  //ö¾ÙBÖÖÀà
			{
				if(num[j]<=B)
				{
					llong cur = 0ll;
					for(int k=0; k<clrpos[j].size(); k++)
					{
						int rb = clrpos[j][k];
						llong tmp = (num[i]-tmp1[rb])*cur%P;
						update(ans2b,tmp);
						update(cur,tmp1[rb]); //lb²»¿ÉµÈÓÚrb, ËùÒÔÏȸüÐÂans2bÔÙ¸üÐÂcur
					}
				}
			}
		}
	}
}

void getans2c() //large-large or small-large ²»ËãAAAA
{
	for(int i=1; i<=m; i++) //ö¾ÙBµÄÖÖÀà
	{
		if(num[i]>B)
		{
			tmp1[0] = 0; for(int j=1; j<=n; j++) tmp1[j] = tmp1[j-1]+(a[j]==i?1:0);
			for(int j=1; j<=m; j++) //ö¾ÙAµÄÖÖÀà
			{
				if(i==j) continue;
				llong cur1 = 0ll,cur2 = 0ll;
				for(int k=0; k<clrpos[j].size(); k++) //ö¾Ùra
				{
					//¸üÐÂans2c
					int ra = clrpos[j][k];
					llong tmp = tmp1[ra]*tmp1[ra]%P*k%P;
					update(ans2c,tmp);
					tmp = tmp1[ra]*(-2ll*cur1-k)%P+P;
					update(ans2c,tmp);
					tmp = cur2+cur1+P;
					update(ans2c,tmp);
					//¸üÐÂcur1,cur2
					update(cur2,tmp1[ra]*tmp1[ra]);
					update(cur1,tmp1[ra]);
				}
			}
		}
	}
	ans2c = ans2c*INV2%P;
}

int main()
{
	scanf("%d",&n); B = sqrt(n)/2;
	for(int i=1; i<=n; i++) scanf("%d",&a[i]),num[a[i]]++,m = max(m,a[i]),clrpos[a[i]].push_back(i);
	for(int i=1; i<=n; i++)
	{
		nxt[i] = lstpos[a[i]];
		lstpos[a[i]] = i;
	}
	getans0();
	getans1();
	getans2a();
	getans2b();
	getans2c();
	ans = ((ans0-ans1-ans2a-ans2b-ans2c)%P+P)%P;
	printf("%lld\n",ans);
	return 0;
}