P3565 [POI2014]HOT-Hotels（樹形dp+長鏈剖分）
P3565 [POI2014]HOT-Hotels
題目大意：
給定一棵樹，在樹上選 3 個點，要求兩兩距離相等，求方案數。
三個點樹上兩兩距離為d存在下面兩種情況
- 某個點三個子樹（保證該點是LCA）中分別由三個點距離它為d
- 對於某一個點，它的 d 級祖先以及子樹內兩個以它為LCA，距它 d 的點
對於情況一，設計dp： f u , j f_{u,j} fu,j表示以 u u u為根的子樹，距i距離為 j j j的點數
對於情況二，設計dp：
g
u
,
j
g_{u,j}
gu,j表示以
u
u
u為根的子樹，兩個點的到LCA距離相等（此出用d表示），LCA到
u
u
對於 f u , j f_{u,j} fu,j的狀態轉移十分顯然： f u , j = ∑ v ∈ s o n u f v , j − 1 f_{u,j}=\sum_{v\in son_u }f_{v,j-1} fu,j=∑v∈sonufv,j−1
而對於 g u , j g_{u,j} gu,j來說存在兩種情況
- 某個子樹內部存在兩個點： g u , j = ∑ v ∈ s o n u g v , j + 1 g_{u,j}=\sum_{v\in son_u}g_{v,j+1} gu,j=∑v∈sonugv,j+1
- 兩個不同的子樹各貢獻一個點，以
u
u
u為LCA： g u , j = ∑ v , w ∈ s o n u , v ≠ w f v , j − 1 × f w , j − 1 g_{u,j}=\sum_{v,w\in son_u,v \ne w} f_{v,j-1}×f_{w,j-1} gu,j=∑v,w∈sonu,v=wfv,j−1×fw,j−1
很明顯，第二中情況
f
v
,
j
−
1
×
f
w
,
j
−
1
,
f
w
,
j
−
1
×
f
v
,
j
−
1
f_{v,j-1}×f_{w,j-1}, f_{w,j-1}×f_{v,j-1}
而對於答案計算來說也存在兩種情況：
首先對於三個點樹上兩兩距離為d的情況都可以看成兩個點在一個子樹中，而另一個點“別處”
- 在子樹 v v v中選一個點，而在其他子樹中選兩個點： g u , j + 1 × f v , j g_{u,j+1}×f_{v,j} gu,j+1×fv,j
- 在子樹 v v v中選兩個點，而在其他子樹中選一個點： f u , j − 1 × g v , j f_{u,j-1}×g_{v,j} fu,j−1×gv,j
同樣為了避免重複計算只需要讓其他子樹搞成
v
v
v前面的子樹即可
注意上述
g
u
,
j
+
1
g_{u,j+1}
gu,j+1和
f
u
,
j
−
1
f_{u,j-1}
fu,j−1都還未算
v
v
v對其的貢獻，這個只需要先計算答案在加貢獻即可。
計算狀態陣列和答案時，都有計算前面子樹的情況，直接運用字首和的思想即可邊計算答案，邊轉移陣列。
這裡要提一點，我們至今沒有提到 g u , 0 g_{u,0} gu,0對答案的貢獻，它確實對答案有貢獻，但是我們已經計算過了，如果在此加上會重複計算。
而其他博主在計算的時候加上 g u , 0 g_{u,0} gu,0實際上增加的時 g w s o n , 1 g_{wson,1} gwson,1即重兒子的貢獻。
根據g的轉移方程可知道： g u , 0 g_{u,0} gu,0的貢獻全部來自於 ∑ v ∈ s o n u g v , 1 \sum_{v\in son_u}g_{v,1} ∑v∈sonugv,1，而計算 f u , j − 1 × g v , j f_{u,j-1}×g_{v,j} fu,j−1×gv,j對答案的貢獻時我們具體寫一下 f u , 0 × g v , 1 f_{u,0}×g_{v,1} fu,0×gv,1而 f u , 0 = 1 f_{u,0}=1 fu,0=1，因此已經計算過 g u , 0 g_{u,0} gu,0的貢獻！！！
然後就是長鏈剖分優化dp，每次儲存長兒子的貢獻，其他兒子暴力合併，每條長鏈都會在鏈頭暴力合併一次時間複雜度 O ( l e n ) O(len) O(len)總的合併時間複雜度 O ( n ) O(n) O(n)
最後如果寫指標版本的話，由於g陣列是倒著轉移的， f f f要多開一倍的記憶體，否則g可能倒回來覆蓋掉 f f f
code
#define IO ios::sync_with_stdio(false);cin.tie(nullptr);cout.tie(nullptr)
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
using ll=long long;
constexpr int N=5010;
int h[N],e[2*N],ne[2*N],idx;
void add(int a,int b){e[idx]=b,ne[idx]=h[a],h[a]=idx++;}
int n;
int dep[N],son[N];
void dfs1(int u,int fa)
{
for(int i=h[u];i!=-1;i=ne[i])
{
int v=e[i];
if(v==fa) continue;
dfs1(v,u);
if(dep[v]>dep[son[u]]) son[u]=v;
}
dep[u]=dep[son[u]]+1;
}
ll pool[4*N];
ll *f[N],*g[N],*now=pool;
ll ans;
void dfs2(int u,int fa)
{
f[u][0]=1;
if(son[u])
{
f[son[u]]=f[u]+1;
g[son[u]]=g[u]-1;
dfs2(son[u],u);
ans+=g[son[u]][1];// 加上重兒子的貢獻
}
for(int i=h[u];i!=-1;i=ne[i])
{
int v=e[i];
if(v==fa||v==son[u]) continue;
f[v]=now;now+=dep[v]<<1;
g[v]=now;now+=dep[v];
dfs2(v,u);
// 邊計算
for(int j=0;j<dep[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<dep[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()
{
IO;
cin>>n;
memset(h,-1,sizeof h);
for(int i=1;i<n;i++)
{
int u,v;
cin>>u>>v;
add(u,v);
add(v,u);
}
dfs1(1,0);
f[1]=now;now+=dep[1]<<1;//多開一倍的記憶體
g[1]=now;now+=dep[1];
dfs2(1,0);
cout<<ans<<'\n';
return 0;
}
菜菜的我搞了一天這個題啊啊啊啊