2019 牛客国庆集训day1 2019 点分治

题目链接：https://ac.nowcoder.com/acm/contest/1099/I

点分治，计算路径数的时候，先将每个点到根的距离模2019，计算的时候就可以O(n)求出数目，对于模2019之后为0的进行特殊处理。

#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;
#define maxn 100005
#define inf 0x3f3f3f3f
int n,k,cnt,root,ans,maxx,head[maxn],size[maxn],son[maxn],vis[maxn],num[maxn];
struct edge{
    int to,next,val;
}e[maxn];
vector<int>dis;
void add(int u,int v,int val)
{
    e[++cnt].to=v;
    e[cnt].next=head[u];
    e[cnt].val=val;
    head[u]=cnt;
}
void dfs_size(int u,int fa)//求各点子树大小 
{
    size[u]=1;son[u]=0;
    for(int i=head[u];i;i=e[i].next)
    {
        int v=e[i].to;
        if(v!=fa&&!vis[v])
        {
            dfs_size(v,u);
            size[u]+=size[v];
            son[u]=max(son[u],size[v]);
        }
    }
}
void dfs_root(int N,int u,int fa)//求重心 
{
    son[u]=max(son[u],N-size[u]);
    if(maxx>son[u])
    {
        root=u;
        maxx=son[u];
    }
    for(int i=head[u];i;i=e[i].next)
    {
        int v=e[i].to;
        if(v!=fa&&!vis[v])
        dfs_root(N,v,u);
    }
}
void dfs_dis(int u,int fa,int val)//求出所有点到根的距离 
{
    val%=2019;
    num[val]++;
    //dis.push_back(val);
    for(int i=head[u];i;i=e[i].next)
    {
        int v=e[i].to;
        if(v!=fa&&!vis[v])
        dfs_dis(v,u,val+e[i].val);
    }
}
int cal(int u,int val)//计算小于等于k的路径数 
{
    for(int i=0;i<=2019;i++)
    num[i]=0;
    //dis.clear();
    dfs_dis(u,0,val);
    //sort(dis.begin(),dis.end());
    int l=0,r=dis.size()-1,ret=0;
    for(int i=1;i<2019;i++)//two-pointer
    {
        if(num[i]&&num[2019-i])ret+=num[i]*num[2019-i];
    }
    ret/=2;
    ret+=num[0]*(num[0]-1)/2;
    return ret;
}
void dfs(int u)
{
    dfs_size(u,0);
    maxx=inf;
    dfs_root(size[u],u,0);
    ans+=cal(root,0);//此时算出的路径数是包括没经过这个根的路径数，后面需要减去这种的路径数 
    vis[root]=1;//将选的roo又被遍历t标记，防止其在之后的分治过程中 
    for(int i=head[root];i;i=e[i].next)
    {
        int v=e[i].to,val=e[i].val;
        if(!vis[v])
        {
            ans-=cal(v,val);//子树所有边加上dis(u,v)后满足的路径数就是需要减去的路径数 
            dfs(v);//递归分治 
        }
    }
}
int main()
{
    while(scanf("%d",&n)!=EOF)
    {
        int u,v,val;
        for(int i=1;i<=n;i++)
        head[i]=vis[i]=0;
        cnt=ans=0;
        for(int i=1;i<n;i++)
        {
            scanf("%d%d%d",&u,&v,&val);
            add(u,v,val);
            add(v,u,val);
        }
        dfs(1);
        printf("%d
",ans);
    }
    return 0;
}