2019牛客国庆集训派对day1 I 题

题目链接: 2019
树形dp搞一下,统计各个儿子间的答案，已经儿子和该节点的答案。统计子节点的时候要记得先减去该子节点的影响，一对点会统计两遍，除2即可

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 2e4 + 5;
const int inf = 0x3f3f3f3f;
const int mod = 2019;
typedef pair < int, int > P;
vector < P > G[N];
int n;
int  cnt[N][2020],ans;
void dfs(int u, int fa) {
    for (int j = 0; j < mod; j++) {cnt[u][j] = 0;}
    for (auto &i : G[u]) {
        int v = i.first, w = i.second;
        if (v == fa) continue;
        dfs(v, u);
        for (int j = 0; j < mod; j++) {
            cnt[u][(j + w) % mod] += cnt[v][j];
        }
    }
    ll res = 0;
        for (auto &i : G[u]) {
            int v = i.first, w = i.second;
            if (v == fa) continue;
            for (int j = 0; j < mod; j++){
                cnt[u][(j + w) % mod] -= cnt[v][j];
            }
            for (int j = 0; j < mod; j++) {
                res += 1LL * cnt[v][j] * (cnt[u][(mod - ((j + w) % mod)) % mod]);
            }
            for (int j = 0; j < mod; j++){
                cnt[u][(j + w) % mod] += cnt[v][j];
            }
        }
    res /= 2;
    ans += res;
    ans += cnt[u][0];
    cnt[u][0]++;
}
int main() {
    int u, v, w;
    while (scanf("%d", &n) != EOF) {
        ans = 0;
        for (int i = 0; i <= n; i++) {
            G[i].clear();
        }
        for (int i = 1; i < n; i++) {
            scanf("%d%d%d", &u, &v, &w);
            G[u].push_back(P(v, w));
            G[v].push_back(P(u, w));
        }
        dfs(1, -1);
        printf("%lld\n",ans);
    }

    return 0;
}