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));
}
}