BZOJ 2152 / Luogu P2634 [国家集训队]聪聪可可 (点分治/树形DP)

题意

一棵树,给定边权,求满足两点之间的路径上权值和为3的倍数的点对数量.

分析

点分治板题,对每个重心求子树下面的到根的距离模3分别为0,1,2的点的个数就行了.
$O(3nlogn)$

CODE

#include<bits/stdc++.h>
using namespace std;
char cb[1<<15],*cs=cb,*ct=cb;
#define getc() (cs==ct&&(ct=(cs=cb)+fread(cb,1,1<<15,stdin),cs==ct)?0:*cs++)
template<class T>inline void read(T &res) {
	char ch; int flg = 1; for(;!isdigit(ch=getc());)if(ch=='-')flg=-flg;
	for(res=ch-'0';isdigit(ch=getc());res=res*10+ch-'0'); res*=flg;
}
const int MAXN = 20005;
int n, fir[MAXN], cnt; bool ban[MAXN];
struct edge{ int to, nxt, w; }e[MAXN<<1];
inline void link(int u, int v, int wt) {
	e[++cnt] = (edge){ v, fir[u], wt }, fir[u] = cnt;
	e[++cnt] = (edge){ u, fir[v], wt }, fir[v] = cnt;
}

int Get_Size(int x, int ff) { //求SIZE
	int re = 1;
	for(int v, i = fir[x]; i; i = e[i].nxt)
		if(!ban[v=e[i].to] && v != ff) re += Get_Size(v, x);
	return re;
}
int Get_Root(int x, int ff, int Size, int &G) { //找重心
	int re = 1; bool flg = true;
	for(int v, i = fir[x]; i; i = e[i].nxt)
		if(!ban[v=e[i].to] && v != ff) {
			int temp = Get_Root(v, x, Size, G);
			if(temp<<1 > Size) flg = false;
			re += temp;
		}
	if(Size-re<<1 > Size) flg = false;
	if(flg) G = x;
	return re;
}
int now[3], tmp[3], Ans;
void Count(int x, int ff, int dis) {
	++tmp[dis];
	for(int v, i = fir[x]; i; i = e[i].nxt)
		if(!ban[v=e[i].to] && v != ff) Count(v, x, (dis+e[i].w)%3);
}
inline void Solve(int x) { //算答案
	now[0] = now[1] = now[2] = 0;
	for(int v, i = fir[x]; i; i = e[i].nxt)
		if(!ban[v=e[i].to]) {
			tmp[0] = tmp[1] = tmp[2] = 0;
			Count(v, x, e[i].w);
			Ans += tmp[0] * now[0] << 1;
			Ans += tmp[1] * now[2] << 1;
			Ans += tmp[2] * now[1] << 1;
			now[0] += tmp[0];
			now[1] += tmp[1];
			now[2] += tmp[2];
		}
	Ans += (now[0]<<1) + 1;
}
void TDC(int x) { //点分治
	int Size = Get_Size(x, 0);
	Get_Root(x, 0, Size, x);
	Solve(x); ban[x] = 1; //打标记
	for(int v, i = fir[x]; i; i = e[i].nxt)
		if(!ban[v=e[i].to]) TDC(v);
}
int main() {
	read(n);
	for(int i = 1, x, y, z; i < n; ++i)
		read(x), read(y), read(z), link(x, y, z%3);
	TDC(1);
	int d = Ans ? __gcd(Ans, n*n) : 1;
	printf("%d/%d\n", Ans/d, n*n/d);
}

$Upd$…sb了…直接树形DP不就完事了… $O(3n)$

CODE

(粘来的代码)

// luogu-judger-enable-o2
#include<cstdio>
#include<cctype>
#define maxn 20005
char cb[1<<15],*cs,*ct;
#define getc() (cs==ct&&(ct=(cs=cb)+fread(cb,1,1<<15,stdin),cs==ct)?0:*cs++)
inline void read(int &a){
    char c;while(!isdigit(c=getc()));
    for(a=c-'0';isdigit(c=getc());a=a*10+c-'0');
}
int n,ans,f[maxn][3],g[3];
int fir[maxn],nxt[maxn<<1],to[maxn<<1],w[maxn<<1],tot;
inline void line(int x,int y,int z){nxt[++tot]=fir[x];fir[x]=tot;to[tot]=y;w[tot]=z;}
void dfs(int u,int pre){
    f[u][0]=1;
    for(int i=fir[u],v;i;i=nxt[i]) if((v=to[i])!=pre){
        dfs(v,u);
        for(int j=0;j<3;j++) g[(j+w[i])%3]=f[v][j];
        for(int j=0;j<3;j++) ans+=f[u][j]*g[j?3-j:0],f[u][j]+=g[j];
    }
}
int gcd(int a,int b){return b?gcd(b,a%b):a;}
int main()
{
    #ifndef ONLINE_JUDGE
        freopen("H.in","r",stdin);
    #endif
    read(n);
    for(int i=1,x,y,z;i<n;i++) read(x),read(y),read(z),line(x,y,z),line(y,x,z);
    dfs(1,0);
    ans=ans*2+n;
    int d=gcd(ans,n*n);
    printf("%d/%d",ans/d,n*n/d);
}