2019牛客国庆集训派对day1

https://ac.nowcoder.com/acm/contest/1099#question

I - 2019

点分治，第一次做的点分治。要先找树的重心，然后分治。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
 
const int INF = 0x3f3f3f3f;
 
const int MAXN = 2e4;
vector<pair<short, short> > G[MAXN + 5];
short dp[MAXN + 5][2019];
short siz[MAXN + 5];
short n, root, maxsiz;
 
ll ans;
 
void dfs1(int u, int p) {
    siz[u] = 1;
    short maxsiz2 = 0;
    for(auto e : G[u]) {
        int v = e.first, w = e.second;
        if(v == p)
            continue;
        dfs1(v, u);
        siz[u] += siz[v];
        maxsiz2 = max(maxsiz2, siz[v]);
    }
    maxsiz2 = max(maxsiz2, (short)(n - siz[u]));
    if(maxsiz2 < maxsiz) {
        maxsiz = maxsiz2;
        root = u;
    }
}
 
void dfs2(int u, int p) {
    dp[u][0] = 1;
    for(auto e : G[u]) {
        int v = e.first, w = e.second;
        if(v == p)
            continue;
        dfs2(v, u);
        for(int k = 0, kw = w; k < 2019; ++k) {
            dp[u][kw] += dp[v][k];
            ++kw;
            if(kw >= 2019)
                kw = 0;
        }
    }
}
 
void calc(int u, int p) {
    ans += dp[u][0] - 1;
    for(auto e : G[u]) {
        int v = e.first, w = e.second;
        if(v == p)
            continue;
        for(int p1 = 0, p2 = (2019 - p1 - w + 2019) % 2019; p1 < 2019; ++p1) {
            ll tmp = 1ll * dp[v][p1] * (dp[u][p2] - dp[v][(p2 - w < 0) ? (2019 + p2 - w) : (p2 - w)]);
            if(tmp >= 1)
                ans += tmp;
            --p2;
            if(p2 < 0)
                p2 += 2019;
        }
        calc(v, u);
    }
}
 
int main() {
#ifdef Yinku
    freopen("Yinku.in", "r", stdin);
#endif // Yinku
    while(~scanf("%d", &n)) {
        for(int i = 1; i <= n - 1; ++i) {
            int u, v, w;
            scanf("%d%d%d", &u, &v, &w);
            G[u].push_back({v, w});
            G[v].push_back({u, w});
        }
 
        root = -1, maxsiz = INF;
        dfs1(1, -1);
        dfs2(root, -1);
 
        ans = 0;
        calc(root, -1);
        printf("%lld
", ans / 2);
 
        for(int i = 0; i <= n; ++i) {
            G[i].clear();
            memset(dp[i], 0, sizeof(dp[i]));
        }
        memset(siz, 0, sizeof(siz[0]) * (n + 1));
    }
}