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]=;son[u]=;
    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%=;
    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=;i<=;i++)
    num[i]=;
    //dis.clear();
    dfs_dis(u,,val);
    //sort(dis.begin(),dis.end());
    int l=,r=dis.size()-,ret=;
    for(int i=;i<;i++)//two-pointer
    {
        if(num[i]&&num[-i])ret+=num[i]*num[-i];
    }
    ret/=;
    ret+=num[]*(num[]-)/;
    return ret;
}
void dfs(int u)
{
    dfs_size(u,);
    maxx=inf;
    dfs_root(size[u],u,);
    ans+=cal(root,);//此时算出的路径数是包括没经过这个根的路径数，后面需要减去这种的路径数
    vis[root]=;//将选的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=;i<=n;i++)
        head[i]=vis[i]=;
        cnt=ans=;
        for(int i=;i<n;i++)
        {
            scanf("%d%d%d",&u,&v,&val);
            add(u,v,val);
            add(v,u,val);
        }
        dfs();
        printf("%d\n",ans);
    }
    return ;
}