洛谷P2634 聪明可可(点分治)

题意：

问树上两点$u,v$路径和为$3$的倍数的有序对$u,v$个数。

分析：

点分治。
将每条到根节点的路径都取模$3$，然后用$dis[0],dis[1],dis[2]$记录个数，
那么一次结果就是$dis[0]\ast dis[0]+2\ast dis[1]\ast dis[2]$。

#include <bits/stdc++.h>
using namespace std;

typedef long long LL;
typedef unsigned long long ULL;
typedef pair<int, int> pii;

const int maxn  = 2e4 + 5;
const int maxm  = 100 + 5;
const int inf   = 0x3f3f3f3f;
const LL  mod   = 1e9 + 7;//19260817
const double pi = acos(-1.0); 

int n, cnt, head[maxn];
int rt, now, tot, siz[maxn], son[maxn], dis[5];
LL ans;
bool vis[maxn];

struct node{
	int v, w, nxt;
}e[maxn << 1];

void addedge(int u, int v, int w){
	e[++cnt] = (node){v, w, head[u]};
	head[u] = cnt; 
}

void getRoot(int x, int f){
	siz[x] = 1, son[x] = 0;
	for(int i = head[x]; i; i = e[i].nxt){
		int v = e[i].v;
		if(v == f || vis[v]) continue;
		getRoot(v, x);
		siz[x] += siz[v];
		son[x] = max(son[x], siz[v]);
	}
	son[x] = max(son[x], now - siz[x]);
	if(son[rt] > son[x] || son[rt] == 0) rt = x;
}

void getDis(int x, int f, int len){
	dis[len % 3]++;
	for(int i = head[x]; i; i = e[i].nxt){
		int v = e[i].v;
		if(v == f || vis[v]) continue;
		getDis(v, x, len + e[i].w); 
	}
}

void solve(int x, int op, int len){
	dis[0] = dis[1] = dis[2] = 0;
	getDis(x, 0, len);
	if(op > 0) ans += 1LL * dis[0] * dis[0] + 2 * dis[1] * dis[2];
	else ans -= (1LL * dis[0] * dis[0] + 2 * dis[1] * dis[2]);
}

void dfs(int x){
	rt = 0, getRoot(x, 0), x = rt;
	vis[x] = 1;
	solve(x, 1, 0);
	for(int i = head[x]; i; i = e[i].nxt){
		int v = e[i].v;
		if(vis[v]) continue;
		solve(v, -1, e[i].w);
		now = siz[v];
		dfs(v);
	}
}

int main(){
	scanf("%d", &n);
	for(int i = 1, u, v, w; i <= n - 1; i++){
		scanf("%d %d %d", &u, &v, &w);
		addedge(u, v, w), addedge(v, u, w);
	}
	now = n;
	getRoot(1, 0);
	dfs(rt);
	int g = __gcd(ans, 1LL * n * n);
	printf("%lld/%lld\n", ans / g, n * n / g);
    return 0;
}