【BZOJ4543】Hotel加強版

題面

bzoj

洛谷

$ps:$在洛谷看題在bzoj交。。。

題解

我們分析一下這個問題，要怎麼樣的點才滿足三點距離兩兩相等呢？

1、存在三個點有共同的$LCA$。

2、存在一個點，使得它到它兩顆不同的子樹種兩點的距離為$d$且它存在$d$級祖先。

考慮$dp:$

設$f[i][j]$表示以$i$為根的子樹中，距離$i$為$j$的點數

$g[i][j]$表示以$i$為根的子樹中兩點到$LCA$距離為$d$,並且它們$LCA$到$i$的距離為$d$的節點數

合併資訊時進行轉移：

$ ans+=g[i][0]\\ ans+=g[i][j]*f[son][j-1]\\ f[i][j]+=f[son][j-1]\\ g[i][j]+=g[son][j+1] $

現在複雜度$O(n^2)$ 注意到這兩個式子

$ f[i][j]+=f[son][j-1]\\ g[i][j]+=g[son][j+1] $

如果我們欽定一個兒子，就不需要再進行重複計算了

我們用指標來描述

$ f[i]=f[son]-1,g[i]=g[son]+1 $

發現鏈上轉移是$O(n)$的

於是我們在樹上做這個事情。

將整棵樹進行長鏈剖分，欽定從重兒子轉移，其他兒子重新計算

是不是和$dsu\;on\;tree$特別像？

那這樣的話，從重兒子轉移$O(1)$,從輕兒子轉移$鏈長O(鏈長)$

這樣總複雜度$O(n)$

程式碼

#include <iostream> 
#include <cstdio> 
#include <cstdlib> 
#include <cstring> 
#include <cmath> 
#include <algorithm> 
using namespace std; 
inline int gi() { 
    register int data = 0, w = 1; 
    register char ch = 0; 
    while (!isdigit(ch) && ch != '-') ch = getchar(); 
    if (ch == '-') w = -1, ch = getchar(); 
    while (isdigit(ch)) data = 10 * data + ch - '0', ch = getchar(); 
    return w * data; 
}
typedef long long ll; 
const int MAX_N = 1e5 + 5; 
struct Graph { int to, next; } e[MAX_N << 1]; int fir[MAX_N], e_cnt = 0; 
void clearGraph() { memset(fir, -1, sizeof(fir)); e_cnt = 0; }
void Add_Edge(int u, int v) { e[e_cnt] = (Graph){v, fir[u]}; fir[u] = e_cnt++; }
int N, dep[MAX_N], son[MAX_N], md[MAX_N]; 
void dfs1(int x, int fa) {
    //md[x] = dep[x] = dep[fa] + 1; 
    for (int i = fir[x]; ~i; i = e[i].next) { 
        int v = e[i].to; if (v == fa) continue; 
        dfs1(v, x);
        if (md[v] > md[son[x]]) son[x] = v, md[x] = md[v]; 
    }
	md[x] = md[son[x]] + 1; 
}
ll *f[MAX_N], *g[MAX_N], tmp[MAX_N << 2], *pos = tmp, ans;
void dfs(int x, int fa) {
    if (son[x]) f[son[x]] = f[x] + 1, g[son[x]] = g[x] - 1, dfs(son[x], x); 
    f[x][0] = 1, ans += g[x][0]; 
    for (int i = fir[x]; ~i; i = e[i].next) {
        int v = e[i].to; if (v == fa || v == son[x]) continue;
        f[v] = pos, pos += md[v] << 1, g[v] = pos, pos += md[v] << 1; 
        dfs(v, x); 
        for (int j = 0; j < md[v]; j++) {
            if (j) ans += f[x][j - 1] * g[v][j]; 
            ans += g[x][j + 1] * f[v][j]; 
        }
        for (int j = 0; j < md[v]; j++) {
            g[x][j + 1] += f[x][j + 1] * f[v][j]; 
            if (j) g[x][j - 1] += g[v][j]; 
            f[x][j + 1] += f[v][j]; 
        } 
    } 
} 
int main () {
    clearGraph(); 
    N = gi();
    for (int i = 1; i < N; i++) {
        int u = gi(), v = gi(); 
        Add_Edge(u, v), Add_Edge(v, u); 
    }
    dfs1(1, 0); f[1] = pos, pos += md[1] << 1, g[1] = pos, pos += md[1] << 1; 
    dfs(1, 0);
    printf("%lld\n", ans); 
    return 0; 
} 
 

玄學問題：去掉第$23$行註釋，並註釋調第$29$行會$RE$