[bzoj2152]聪聪可可——点分治

Brief Descirption

给定一棵带权树，您需要统计路径长度为3的倍数的路径长度

Algorithm Analyse

点分治。
考察经过重心的路径。统计出所有deep，统计即可。

Code

#include <cstdio>
#include <cstring>
#include <vector>
#define ll long long
const int maxn = 40005;
int n, ans, rt, sum;
using std::max;
struct edge {
  int to, weigh;
};
std::vector<edge> G[maxn];
int vis[maxn], size[maxn], f[maxn], cnt[maxn], deep[maxn];
void add_edge(int from, int to, int x) {
  G[from].push_back((edge){to, x});
  G[to].push_back((edge){from, x});
}
int gcd(int a, int b) { return b == 0 ? a : gcd(b, a % b); }
void getroot(int x, int fa) {
  size[x] = 1;
  f[x] = 0;
  for (int i = 0; i < G[x].size(); i++) {
    edge &e = G[x][i];
    if (!vis[e.to] && e.to != fa) {
      getroot(e.to, x);
      size[x] += size[e.to];
      f[x] = max(f[x], size[e.to]);
    }
  }
  f[x] = max(f[x], sum - size[x]);
  if (f[x] < f[rt])
    rt = x;
}
void getdeep(int x, int fa) {
  cnt[deep[x]]++;
  for (int i = 0; i < G[x].size(); i++) {
    edge &e = G[x][i];
    if (!vis[e.to] && e.to != fa) {
      deep[e.to] = (deep[x] + e.weigh) % 3;
      getdeep(e.to, x);
    }
  }
}
int cal(int x, int now) {
  cnt[0] = cnt[1] = cnt[2] = 0;
  deep[x] = now;
  getdeep(x, 0);
  return cnt[1] * cnt[2] * 2 + cnt[0] * cnt[0];
}
void work(int x) {
  ans += cal(x, 0); //统计不同子树通过重心的个数
  vis[x] = 1;
#ifndef ONLINE_JUDGE
  printf("In root %d: %d
", rt, ans);
#endif
  for (int i = 0; i < G[x].size(); i++) {
    edge &e = G[x][i];
    if (!vis[e.to]) {
      ans -= cal(e.to, e.weigh); //去除在同一个子树中被重复统计的
      rt = 0;
      sum = size[e.to];
      getroot(e.to, 0);
      work(rt); // Decrease and Conquer
    }
  }
}
int main() {
#ifndef ONLINE_JUDGE
  freopen("input", "r", stdin);
#endif
  scanf("%d", &n);
  for (int i = 1; i < n; i++) {
    int u, v, b;
    scanf("%d %d %d", &u, &v, &b);
    add_edge(u, v, b % 3);
  }
  memset(vis, 0, sizeof(vis));
  rt = 0;
  f[0] = sum = n;
  ans = 0;
  getroot(1, 0);
  work(rt);
  int t = gcd(ans, n * n);
  printf("%d/%d
", ans / t, n * n / t);
  return 0;
}
/*
5
1 2 1
1 3 2
1 4 1
2 5 3
*/