【点分治】【FFT】CDOJ1562 Amaz1ng Prime

统计路径的时候，显然用母函数的思想，可以用FFT来方便统计。

注意！要减去路径两个端点相同的情况！然后再除以二！这样防止重复。

还有就是说啊，点分治的正确姿势还是应该用所有子树的答案减去各个子树分别的答案。否则复杂度好像是不太对哈。

除非……有一种情况是这样的……才能让每次直接统计当前子树和之前所有子树的答案而不重复统计再减去的写法复杂度正确。

即：用个set或者哈希表之类的东西存储之前所有子树的路径，然后我只枚举当前子树的路径，往set或者哈希表里面去查询。查询完了以后，再把当前子树塞到set/哈希表里。

这样复杂度正确的原因是，set和哈希表不会因为其储存元素的多少而使复杂度快速上升。

FFT貌似是得开4倍数组哈。

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<iostream>
#include<cmath>
using namespace std;
typedef long long ll;
#define MAXN 100010
bool notPrime[200010];
void Shai()
{
    notPrime[1]=1;
    notPrime[0]=1;
    for(ll i=2;i<=200000ll;i++)
      for(ll j=i*i;j<=200000ll;j+=i)
        notPrime[j]=1;
}
const double PI = acos(-1.0);
struct Complex{
	double real,image;
	Complex(double _real,double _image){
		real=_real;
		image=_image;
	}
	Complex(){}
};
Complex operator + (const Complex &c1,const Complex &c2){
	return Complex(c1.real+c2.real,c1.image+c2.image);
}
Complex operator - (const Complex &c1,const Complex &c2){
	return Complex(c1.real-c2.real,c1.image-c2.image);
}
Complex operator * (const Complex &c1,const Complex &c2){
	return Complex(c1.real*c2.real-c1.image*c2.image,c1.real*c2.image+c1.image*c2.real);
}
int rev(int id,int len){
    int ret=0;
    for(int i=0;(1<<i)<len;++i){
		ret<<=1;
		if(id&(1<<i)){
			ret|=1;
		}
	}
	return ret;
}
Complex tmp[200010];
//当DFT==1时是DFT, DFT==-1时则是逆DFT
void IterativeFFT(Complex A[],int len, int DFT){//对长度为len(2的幂)的数组进行DFT变换
	for(int i=0;i<len;++i){
		tmp[rev(i,len)]=A[i];
	}
	for(int i=0;i<len;++i){
		A[i]=tmp[i];
	}
	for(int s=1;(1<<s)<=len;++s){
		int m=(1<<s);
		Complex wm=Complex(cos(DFT*2*PI/m),sin(DFT*2*PI/m));
		for(int k=0;k<len;k+=m){//这一层结点包含的数组元素个数都是(1<<s)
			Complex w=Complex(1,0);
			for(int j=0;j<(m>>1);++j){//折半引理，根据两个子节点计算父节点
				Complex t=w*A[k+j+(m>>1)];
				Complex u=A[k+j];
				A[k+j]=u+t;
				A[k+j+(m>>1)]=u-t;
				w=w*wm;
			}
		}
	}
	if(DFT==-1){
		for(int i=0;i<len;++i){
			A[i].real/=len;
			A[i].image/=len;
		}
	}
}
typedef pair<int,int> Point;
int n;
ll ans;
int v[MAXN<<1],w[MAXN<<1],first[MAXN],__next[MAXN<<1],en;
void AddEdge(const int &U,const int &V,const int &W)
{
    v[++en]=V;
    w[en]=W;
    __next[en]=first[U];
    first[U]=en;
}
bool centroid[MAXN];//顶点是否已经作为重心删除的标记
int size[MAXN];//以该顶点为根的子树的大小
//计算子树的大小
int calc_sizes(int U,int Fa)
{
    int res=1;
    for(int i=first[U];i;i=__next[i])
      if(v[i]!=Fa&&(!centroid[v[i]]))
        res+=calc_sizes(v[i],U);
    return size[U]=res;
}
//查找重心的递归函数，nn是整个子树的大小
//在以U为根的子树中寻找一个顶点，使得删除该顶点后得到的最大子树的顶点数最少
//返回值为(最大子树的顶点数,顶点编号)
Point calc_centroid(int U,int Fa,int nn)
{
    Point res=make_pair(2147483647,-1);
    int sum=1,maxv=0;
    for(int i=first[U];i;i=__next[i])
      if(v[i]!=Fa&&(!centroid[v[i]]))
        {
          res=min(res,calc_centroid(v[i],U,nn));
          maxv=max(maxv,size[v[i]]);
          sum+=size[v[i]];
        }
    maxv=max(maxv,nn-sum);
    res=min(res,make_pair(maxv,U));
    return res;
}
int td[MAXN],en2,ds[MAXN],en3;
//计算子树中所有顶点到重心的距离的递归函数
void calc_dis(int U,int Fa,int d)
{
    td[en2++]=d;
    for(int i=first[U];i;i=__next[i])
      if(v[i]!=Fa&&(!centroid[v[i]]))
        calc_dis(v[i],U,d+w[i]);
}
Complex fft[200100];
ll calc_pairs(int dis[],int En)
{
	int lim=0;
	for(int i=0;i<En;++i){
		lim=max(lim,dis[i]);
	}
	++lim;
	int len;
	for(int i=0;;++i){
		if((1<<(i-1))>=lim){
			len=(1<<i);
			break;
		}
	}
	for(int i=0;i<len;++i){
		fft[i]=Complex(0,0);
	}
	for(int i=0;i<En;++i){
		fft[dis[i]].real+=1.0;
	}
	IterativeFFT(fft,len,1);
	for(int i=0;i<len;++i){
		fft[i]=fft[i]*fft[i];
	}
	IterativeFFT(fft,len,-1);
    ll res=0;
	for(int i=0;i<len;++i){
		if(!notPrime[i]){
			res+=(ll)(fft[i].real+0.5);
		}
	}
	for(int i=0;i<En;++i){
		if(!notPrime[dis[i]<<1]){
			--res;
		}
	}
	return (res>>1);
}
void solve(int U)
{
    calc_sizes(U,-1);
    int s=calc_centroid(U,-1,size[U]).second;
    centroid[s]=1;
    //情况1:递归统计按重心s分割后的子树中的对数
    for(int i=first[s];i;i=__next[i])
      if(!centroid[v[i]])
        solve(v[i]);
    //情况2:统计经过重心s的对数
    en3=0; ds[en3++]=0;
    for(int i=first[s];i;i=__next[i])
      if(!centroid[v[i]])
        {
          en2=0; calc_dis(v[i],s,w[i]);
          ans-=calc_pairs(td,en2);//先把重复统计的部分(即情况1)减掉
          memcpy(ds+en3,td,en2*sizeof(int)); en3+=en2;
        }
    ans+=calc_pairs(ds,en3);
    centroid[s]=0;
}
int main()
{
//	freopen("cdoj1562.in","r",stdin);
	Shai();
    int a,b,c;
    scanf("%d",&n);
    for(int i=1;i<n;++i)
      {
        scanf("%d%d%d",&a,&b,&c);
        AddEdge(a,b,c);
		AddEdge(b,a,c);
      }
    solve(1);
    cout<<ans<<endl;
    return 0;
}