题目大意
给定一棵 (n(nleq 5000)) 个点的树,求选出三个不同的点,使得三个点两两之间间距相等的方案数。
题解
设 (f[u][j]) 表示以 (u) 为根的子树中,到 (u) 距离为 (j) 的结点的数目,设 (g[v][j]) 表示lca在以 (v) 为根的子树中 ,且到 (v) 的距离相同,且 (v) 到 (u) 的距离加 (j) 等于它们到 (v) 的距离的结点对数。
每访问到 (u) 的一个儿子 (v),就有
[ans+=f[u][j] imes g[v][j+1]+f[v][j-1] imes g[u][j]\
g[u][j+1]+=f[u][j+1] imes f[v][j]\
g[u][j-1]+=g[v][j]\
f[u][j]+=f[v][j-1]
]
时间复杂度 (O(n^2))。
Code
#include <bits/stdc++.h>
using namespace std;
#define RG register int
#define LL long long
template<typename elemType>
inline void Read(elemType& T) {
elemType X = 0, w = 0; char ch = 0;
while (!isdigit(ch)) { w |= ch == '-';ch = getchar(); }
while (isdigit(ch)) X = (X << 3) + (X << 1) + (ch ^ 48), ch = getchar();
T = (w ? -X : X);
}
struct Graph {
struct edge { int Next, to; };
edge G[10010];
int head[5010];
int cnt;
Graph() :cnt(2) {}
void clear(int n) {
cnt = 2;fill(head, head + n + 2, 0);
}
void add_edge(int u, int v) {
G[cnt].to = v;
G[cnt].Next = head[u];
head[u] = cnt++;
}
};
Graph G;
int N;
int Height[5005];
LL f[5005][5005], g[5005][5005], ans;
void DFS(int u, int fa) {
Height[u] = 0;
f[u][0] = 1;
for (int i = G.head[u];i;i = G.G[i].Next) {
int v = G.G[i].to;
if (v == fa) continue;
DFS(v, u);
Height[u] = max(Height[u], Height[v] + 1);
for (int j = 0;j <= Height[u];++j)
ans += f[u][j] * g[v][j + 1] + g[u][j] * f[v][j - 1];
for (int j = 0;j <= Height[u];++j) {
g[u][j + 1] += f[u][j + 1] * f[v][j];
if (j >= 1) {
g[u][j - 1] += g[v][j];
f[u][j] += f[v][j - 1];
}
}
}
}
int main() {
Read(N);
for (int i = 1;i <= N - 1;++i) {
int u, v;
Read(u);Read(v);
G.add_edge(u, v);
G.add_edge(v, u);
}
DFS(1, 0);
printf("%lld
", ans);
return 0;
}
长链剖分优化dp
BZOJ 4543这道题 (n) 扩大到了 (10^5)。
#include <bits/stdc++.h>
using namespace std;
#define RG register int
#define LL long long
template<typename elemType>
inline void Read(elemType& T) {
elemType X = 0, w = 0; char ch = 0;
while (!isdigit(ch)) { w |= ch == '-';ch = getchar(); }
while (isdigit(ch)) X = (X << 3) + (X << 1) + (ch ^ 48), ch = getchar();
T = (w ? -X : X);
}
const int maxn = 100005;
struct Graph {
struct edge { int Next, to; };
edge G[maxn << 1];
int head[maxn];
int cnt;
Graph() :cnt(2) {}
void clear(int n) {
cnt = 2;fill(head, head + n + 2, 0);
}
void add_edge(int u, int v) {
G[cnt].to = v;
G[cnt].Next = head[u];
head[u] = cnt++;
}
};
Graph G;
int N;
int Height[maxn], Hson[maxn];
LL temp[maxn * 5], * f[maxn], * g[maxn], * id = temp, ans;
void DFS_Init(int u, int fa) {
Height[u] = 1;
for (int i = G.head[u];i;i = G.G[i].Next) {
int v = G.G[i].to;
if (v == fa) continue;
DFS_Init(v, u);
Height[u] = max(Height[u], Height[v] + 1);
if (Height[v] > Height[Hson[u]]) Hson[u] = v;
}
return;
}
void DFS(int u, int fa) {
if (Hson[u]) {
f[Hson[u]] = f[u] + 1;
g[Hson[u]] = g[u] - 1;
DFS(Hson[u], u);
}
f[u][0] = 1;ans += g[u][0];
for (int i = G.head[u];i;i = G.G[i].Next) {
int v = G.G[i].to;
if (v == fa || v == Hson[u]) continue;
f[v] = id;id += Height[v] << 1;
g[v] = id;id += Height[v] << 1;
DFS(v, u);
for (int j = 0;j < Height[v];++j) {
if (j) ans += g[v][j] * f[u][j - 1];
ans += g[u][j + 1] * f[v][j];
}
for (int j = 0;j < Height[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() {
Read(N);
for (int i = 1;i <= N - 1;++i) {
int u, v;
Read(u);Read(v);
G.add_edge(u, v);
G.add_edge(v, u);
}
DFS_Init(1, 0);
f[1] = id;id += Height[1] << 1;
g[1] = id;id += Height[1] << 1;
DFS(1, 0);
printf("%lld
", ans);
return 0;
}