C++11(clang++ 3.9) 解法, 执行用时: 229ms, 内存消耗: 9184K, 提交时间: 2020-05-10 14:03:17

#include<iostream>
#include<algorithm>
#include<vector>
#include<cstdio>
using namespace std;
#define maxn 100005
#define inf 0x3f3f3f3f
#define ll long long
ll n,k,cnt,root,ans,maxx,head[maxn],size[maxn],son[maxn],vis[maxn],num[maxn];
struct edge{
    ll to,next,val;
}e[maxn*4];

vector<ll>dis;

void add(ll u,ll v,ll val)
{
    e[++cnt].to=v;
    e[cnt].next=head[u];
    e[cnt].val=val;
    head[u]=cnt;
}
void dfs_size(ll u,ll fa)//求各点子树大小
{
    size[u]=1;son[u]=0;
    for(ll i=head[u];i;i=e[i].next)
    {
        ll 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(ll N,ll u,ll fa)//求重心
{
    son[u]=max(son[u],N-size[u]);
    if(maxx>son[u])
    {
        root=u;
        maxx=son[u];
    }
    for(ll i=head[u];i;i=e[i].next)
    {
        ll v=e[i].to;
        if(v!=fa&&!vis[v])
        dfs_root(N,v,u);
    }
}
void dfs_dis(ll u,ll fa,ll val)//求出所有点到根的距离
{
    val%=2020;
    num[val]++;
    //dis.push_back(val);
    for(ll i=head[u];i;i=e[i].next)
    {
        ll v=e[i].to;
        if(v!=fa&&!vis[v])
        dfs_dis(v,u,val+e[i].val);
    }
}
ll cal(ll u,ll val)//计算小于等于k的路径数
{
    for(ll i=0;i<2020;i++)
    num[i]=0;
    //dis.clear();
    dfs_dis(u,0,val);
    //sort(dis.begin(),dis.end());
    ll ret=0;
    for(ll i=1;i<2020;i++)//two-pointer
    {
        if(num[i]&&num[2020-i])ret+=num[i]*num[2020-i];
    }
    ret/=2;
    ret+=num[0]*(num[0]-1)/2;

    return ret;
}
void dfs(ll u)
{
    dfs_size(u,0);
    maxx=inf;
    dfs_root(size[u],u,0);
    ans+=cal(root,0);//此时算出的路径数是包括没经过这个根的路径数，后面需要减去这种的路径数
    vis[root]=1;//将选的roo又被遍历t标记，防止其在之后的分治过程中
    for(ll i=head[root];i;i=e[i].next)
    {
        ll v=e[i].to,val=e[i].val;
        if(!vis[v])
        {
            ans-=cal(v,val);//子树所有边加上dis(u,v)后满足的路径数就是需要减去的路径数
            dfs(v);//递归分治
        }
    }
}

int main()
{
    while(scanf("%lld",&n)!=EOF)
    {
        ll u,v,val;
        for(ll i=1;i<=n;i++)
        head[i]=vis[i]=0;
        cnt=ans=0;
        for(ll i=1;i<n;i++)
        {
            scanf("%lld%lld%lld",&u,&v,&val);
            add(u,v,val);
            add(v,u,val);
        }
        dfs(1);
        printf("%lld\n",ans);
    }
    return 0;
}

C++14(g++5.4) 解法, 执行用时: 90ms, 内存消耗: 13264K, 提交时间: 2020-05-10 14:32:16

#include<bits/stdc++.h>

typedef long long ll;

using namespace std;
#pragma GCC optimize(2)
#define Init1 ios::sync_with_stdio(false)
#define Init2 cin.tie(0)
#define INF 0x3f3f3f3f
#define ALL(x) x.begin(), x.end()

const int N = 2e5 + 105;
vector<pair< int, int > >edge[N];

int dfn[N], dist[N], sz[N], dfncnt = 0, rt, sum, vis[N], mx[N];
void calcsiz(int u, int fa){
    sz[u] = 1;
    mx[u] = 0;
    int v;
    for(auto it : edge[u]){
        v = it.first;
        if(v == fa || vis[v]) continue;
        calcsiz(v, u);
        sz[u] += sz[v];
        mx[u] = max(mx[u], sz[v]);
    }
    mx[u] = max(mx[u], sum - sz[u]);
    if(mx[rt] > mx[u]) rt = u;
}
int tf[2050], cnt[N], tot = 0;
queue<int>que;
ll ans = 0;
void calcdist(int u, int fa){
    cnt[++tot] = dist[u];
    int v, val;
    for(auto it : edge[u]){
        v = it.first, val = it.second;
        if(vis[v] || v == fa) continue;
        dist[u] = (dist[v] + val)%2020, calcdist(v, u);
    }
}
void dfz(int u, int fa){
    vis[u] = 1;
    int v, val, p;
    tf[0] = 1; que.push(0);
    for(auto it : edge[u]){
        v = it.first, val = it.second;
        if(vis[v] || v == fa) continue;
        dist[v] = val%2020;
        calcdist(v, u);
        for(int i = 1; i <= tot; ++i){
            p = 2020 - cnt[i];
            if(p == 2020) p = 0;
            ans += tf[p];
        }
        for(int i = 1; i <= tot; ++i){
            tf[cnt[i]]++; que.push(cnt[i]);
        }
        tot = 0;
    }
    while(!que.empty()){
        tf[que.front()] = 0; que.pop();
    }
    for(auto it : edge[u]){
        v = it.first, val = it.second;
        if(vis[v] || v == fa) continue;
        rt = 0, mx[rt] = INF; sum = sz[v];
        calcsiz(v, u);
        calcsiz(rt, 0);
        dfz(rt, 0);
    }
}
int main(){
    Init1; Init2;
    int n;
    int u, v, val;
    scanf("%d", &n);
    for(int i = 1; i < n; ++i){
        scanf("%d%d%d", &u, &v, &val);
        edge[u].push_back(make_pair(v, val));
        edge[v].push_back(make_pair(u, val));
    }
    rt = 0, mx[rt] = INF; sum = n;
    calcsiz(1, 0);
    calcsiz(rt, 0);
    dfz(rt, 0);
    printf("%lld\n", ans);
    // system("pause");
    return 0;

}
//     1 2 0
// 2 3 0
// 3 4 0
// 2 5 0