[POI2014]HOT-Hotels dp + 长链剖分

[POI2014]HOT-Hotels

写出最暴力的转移dp之后， 可以发现用长链剖分优化。

#include
#define LL long long
using namespace std;

const int N = (int)2e6 + 7;

int n, len[N], son[N];
vector<int> G[N];

LL dp[N];
LL *f[N], *g[N];
LL *id = dp;

void gao(int u, int fa) {
    son[u] = len[u] = 0;
    for(auto &v : G[u]) {
        if(v == fa) continue;
        gao(v, u);
        if(len[v] > len[u]) {
            son[u] = v;
            len[u] = len[v];
        }
    }
    len[u]++;
}

LL ans = 0;

void dfs(int u, int fa) {
    f[u][0] = 1;
    if(son[u]) {
        f[son[u]] = f[u] + 1;
        g[son[u]] = g[u] - 1;
        dfs(son[u], u);
    }
    ans += g[u][0];
    for(auto &v : G[u]) {
        if(v == fa || v == son[u]) continue;
        f[v] = id; id += len[v] * 2;
        g[v] = id; id += len[v] * 2;
        dfs(v, u);
        for(int i = 1; i <= len[v]; i++) {
            ans += g[u][i] * f[v][i - 1];
            if(i >= 2) ans += f[u][i - 2] * g[v][i - 1];
        }
        for(int i = 1; i <= len[v]; i++) {
            if(i >= 2) g[u][i - 2] += g[v][i - 1];
            g[u][i] += f[u][i] * f[v][i - 1];
        }
        for(int i = 1; i <= len[v]; i++) {
            f[u][i] += f[v][i - 1];
        }
    }
}

int main() {
    scanf("%d", &n);
    for(int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    gao(1, 0);
    f[1] = id; id += len[1] * 2;
    g[1] = id; id += len[1] * 2;
    dfs(1, 0);
    printf("%lld\n", ans);
    return 0;
}