题意
求一棵树上,两两距离相等的三个点的三元组(无序)的个数。
题解
CODE
代码中的 f , g f,g f,g对应题解中的 n u m , w a y num,way num,way
#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
inline void rd(int &x) {
char ch; int flg=1; while(!isdigit(ch=getchar()))if(ch=='-')flg=-flg;
x = 0; do x=x*10+ch-'0'; while(isdigit(ch=getchar())); x*=flg;
}
const int MAXN = 300005;
int n, m, fir[MAXN], to[MAXN<<1], nxt[MAXN<<1], cnt;
inline void link(int u, int v) {
to[++cnt] = v, nxt[cnt] = fir[u], fir[u] = cnt;
to[++cnt] = u, nxt[cnt] = fir[v], fir[v] = cnt;
}
int mxd[MAXN], son[MAXN];
void dfs1(int u, int ff) {
for(int i = fir[u], v; i; i = nxt[i])
if((v=to[i]) != ff) {
dfs1(v, u);
if(mxd[v] > mxd[son[u]]) son[u] = v;
}
mxd[u] = mxd[son[u]] + 1;
}
LL tmp[MAXN<<2], *f[MAXN], *g[MAXN], *id=tmp; LL ans;
void dfs(int u, int ff) {
if(son[u]) f[son[u]] = f[u] + 1, g[son[u]] = g[u] - 1, dfs(son[u], u);
f[u][0] = 1; ans += g[u][0];
for(int i = fir[u], v; i; i = nxt[i])
if((v=to[i]) != ff && v != son[u]) {
f[v] = id; id += mxd[v]<<1;
g[v] = id; id += mxd[v]<<1;
dfs(v, u);
for(int j = 0; j < mxd[v]; ++j) {
if(j) ans += f[u][j-1] * g[v][j];
ans += g[u][j+1] * f[v][j];
}
for(int j = 0; j < mxd[v]; ++j) {
g[u][j+1] += f[u][j+1] * f[v][j];
if(j) g[u][j-1] += g[v][j];
f[u][j+1] += f[v][j];
}
}
}
int main () {
rd(n);
for(int i = 1, x, y; i < n; ++i)
rd(x), rd(y), link(x, y);
dfs1(1, 0);
f[1] = id; id += mxd[1]<<1;
g[1] = id; id += mxd[1]<<1;
dfs(1, 0);
printf("%lld\n", ans);
}