2019牛客国庆集训派对day1-2019(I)

题目连接

图片

核心：

感受又是一个树形dp. dp[rt][k]：表示以rt为根的树中，到rt距离为%2019=k的个数。 更新方式：来了一颗新树x, rt这边到rt距离为d1的点，与x这边到x距离为d2的点，若(d1+d2+val)%2019==0, 则ans+=dp[rt][d1]*dp[val][d2],val表示rt与x之间的距离; 以后更新dp[rt][d], dp[rt][(d+val)%2019]+=dp[rt][d]

总结：

先求ans, 再更新dp[rt][d]

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
struct node {
    int to;
    int val;
};
const int N=2e4+3;
const int mod=2019;
int d[N][mod+3];
bool vis[N];
vector<node> g[N];
int n; LL ans;
void dfs(int rt) {
    d[rt][0]+=1; // 重要
    vis[rt]=1;
    for (int i=0;i<g[rt].size();i++) {
        int nxt=g[rt][i].to;
        int val=g[rt][i].val;
        if(vis[nxt]) continue;
        dfs(nxt);
        for(int k=0;k<mod;k++) {
            int x=(mod*2-k-val)%mod;
            ans+=d[nxt][x]*d[rt][k]; // 统一全部的距离是通过rt与nxt的一条路径
        }
        for(int k=0;k<mod;k++)
            d[rt][(k+val)%mod]+=d[nxt][k];
    }
    // 形式二 d[rt][0]+=1;
    //       ans+=d[rt][0]-1;
}
int main()
{
    while(~scanf("%d",&n)) {
        ans=0;
        memset(d,0,sizeof(d));
        memset(vis,0,sizeof(vis));
        for (int i=0;i<N;i++)
            g[i].clear();
        for (int i=1;i<n;i++) {
            int a,b,c;
            scanf("%d %d %d",&a,&b,&c);
            // 个人傻逼错误
            // if(a<b) swap(a,b); 这样会改变图的形状
            node tmp={b,c}; g[a].push_back(tmp);
            tmp.to=a;       g[b].push_back(tmp);
        }
        dfs(1);
        printf("%lld\n",ans);
    }
    return 0;

}