Codechef June Challenge 2014 #Sereja and Arcs -- 容斥 + 阈值 + 树状数组

传送门

题目大意

数轴上有 $n$ 个点，给你每个点的坐标和颜色。相同颜色的每对点$A(x,0),B(y,0)$会产生一个直径为$AB$的圆。求有多少对颜色不同的圆相交。

容斥。将问题转化为求有多少对颜色不同的圆不相交。
将相同颜色的点用同一个字母表示，那么只有$AABB$和$ABBA$是满足条件的。

考虑$AABB$的情况：
从左到右枚举每个点作为右端点。$O(n)$。

考虑$ABBA$的情况：
阈值。将出现次数大于 $S$ 的颜色记为 $P$ ，小于 $S$ 的颜色记为 $Q$ 。
那么共有4种情况：$QPPQ~,~P_1 P_2 P_2 P_1~,~PQQP~,~Q_1Q_2Q_2Q_1$。

• $QPPQ~,~P_1 P_2 P_2 P_1$的情况：对于每个 $P$，枚举每种颜色，求出这种颜色出现的每个位置之前它出现的次数 $p_i$ 。枚举它出现的每个位置，第 $i$ 个位置对答案的贡献就是$\sum_{j=1}^{i-1} {(p_i-p_j)*(p_i-p_j-1)\over 2}$，维护 $\sum_{j-1}^{i-1}p_i$ 和 $\sum_{j-1}^{i-1}p_i^2$ 就可以$O(1)$求出了。由于 $P$ 最多只有$n\over S$个，时间复杂度为$O({n^2\over S})$。
• $PQQP$的情况：对于每个 $P$，枚举每个 $Q$ ，求出这种颜色出现的每个位置之前它出现的次数 $p_i$ 和之后它出现的次数 $s_i$ ，第 $i$ 个位置对答案的贡献就是$\sum_{j=1}^{i-1} s_i*p_j$。同样可以$O(1)$求出。时间复杂度为$O({n^2\over S})$。
• $Q_1Q_2Q_2Q_1$的情况：类似二维数点。从左到右枚举每个点 $i$ 作为右端点，再枚举前面颜色与它相同的点 $j$ 作为左端点，那么贡献就是$j\le L\le R\le i$的$L,R$对数。由于之前的 $R$ 都小于 $i$ ，只需要用树状数组维护$j\le L$的 $L$ 个数就可以了。时间复杂度$O(nS\log n)$。

取$S=\sqrt{n\over{\log n}}$，时间复杂度为$O(n\sqrt{n\log n})$。
由于树状数组常数小， S <script type="math/tex" id="MathJax-Element-46">S</script> 可以取得更大。

代码

#include<cmath>
#include<cstdio>
#include<vector>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
#define N 100010
#define ll long long
#define I 500000004
#define M 1000000007
vector<int>g[N];
int Ans;
int s1,s2,Sum,S,i,j,k,n,m,p,s[N],c[N],a[N],w[N],l[N],f[N],x,Res,t[N],F[N];
inline ll Query(int x){
    int Ans=0;
    for(x=n-x+1;x;x-=x&-x)Ans+=t[x];
    return Ans;
}
inline void Update(int x){
    for(x=n-x+1;x<=n;x+=x&-x)t[x]++;
}
int main(){
    scanf("%d",&n);
    for(i=1;i<=n;i++)scanf("%d",&a[i]),w[i]=++l[a[i]],g[a[i]].push_back(i),c[a[i]]++,m=max(m,a[i]);
    for(i=1;i<=m;i++)Sum+=1ll*c[i]*(c[i]-1)/2%M;
    for(i=1;i<=n;i++)
    Ans=(Ans+1ll*(Sum-1ll*(c[a[i]]-w[i]+1)*(c[a[i]]-w[i])/2%M)*(w[i]-1)%M)%M,Sum=(Sum-c[a[i]]+w[i])%M;
    S=sqrt(n);
    for(i=1;i<=m;i++)
    if(c[i]>S){
        memset(f,0,sizeof(f));
        memset(F,0,sizeof(F));
        for(j=0;j<g[i].size();j++)f[g[i][j]]=F[g[i][j]]=1;
        for(j=2;j<=n;j++)f[j]+=f[j-1];
        for(j=n-1;j;j--)F[j]+=F[j+1];
        for(j=1;j<=m;j++)
        if(j!=i){
            s1=s2=0;
            for(k=0;k<g[j].size();k++){
                x=f[g[j][k]];
                Res=1ll*k*(1ll*x*x%M-x)%M;
                Res=(Res+s2)%M;
                Res=(Res+1ll*s1*(1-2*x)%M)%M;
                Res=1ll*Res*I%M;
                Ans=(Ans+Res)%M;
                s1=(s1+x)%M;s2=(s2+1ll*x*x%M)%M;
            }
            if(c[j]<=S){
                s1=0;
                for(k=(int)g[j].size()-1;k>=0;k--){
                    Ans=(Ans+1ll*f[g[j][k]]*s1%M);
                    s1=(s1+F[g[j][k]]);
                }
            }
        }
    }
    for(i=1;i<=n;i++)
    if(c[a[i]]<=S){
        for(j=0;j<w[i]-1;j++){
            Ans=(Ans+Query(g[a[i]][j]))%M;
            Ans=(Ans-1ll*(w[i]-j-1)*(w[i]-j-2)/2%M)%M;
            Update(g[a[i]][j]);
        }
    }
    k=0;
    for(p=0,i=1;i<=m;i++)if(c[i]>1)p=(p+1ll*c[i]*(c[i]-1)/2%M)%M,k++;
    Res=0;
    for(i=1;i<=m;i++)Res=(Res+1ll*c[i]*(c[i]-1)/2%M*(p-1ll*c[i]*(c[i]-1)/2%M)%M)%M;
    Res=1ll*Res*I%M;
    if(k<=1)puts("0");else cout<<((Res-Ans)%M+M)%M<<endl;
    return 0;
}