4543: [POI2014]Hotel加强版

4543: [POI2014]Hotel加强版

链接

分析：

　　f[u][i]表示子树u内，距离u为i的点的个数，g[u][i]表示在子树u内，已经选了两个深度一样的点，还需要在距离u为i的一个点作为第三个点。

　　然后就可以利用这两个数组统计答案了。

　　ans+=g[u][j]*f[v][j-1]+f[u][j]*g[v][j+1];

　　如果直接合并f和g，复杂度是$O(n^2)$的，如果可以启发式合并，复杂度是$O(nlogn)$的，如果是长链剖分，复杂度是$O(n)$的。

　　长链剖分就是按照子树内最长的链就行剖分，根节点保留深度最大的那个点的信息，其他的点暴力合并。

代码：

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<iostream>
#include<cmath>
#include<cctype>
#include<set>
#include<queue>
#include<vector>
#include<map>
using namespace std;
typedef long long LL;

inline int read() {
    int x=0,f=1;char ch=getchar();for(;!isdigit(ch);ch=getchar())if(ch=='-')f=-1;
    for(;isdigit(ch);ch=getchar())x=x*10+ch-'0';return x*f;
}

const int N = 100005;
struct Edge{ int to, nxt; } e[N << 1]; 
int head[N], len[N], fa[N], son[N], En;
LL ans, *f[N], *g[N], ft[N], gt[N << 1], *fp = ft, *gp = gt;

inline void add_edge(int u,int v) {
    ++En; e[En].to = v, e[En].nxt = head[u]; head[u] = En;
    ++En; e[En].to = u, e[En].nxt = head[v]; head[v] = En;
}
void dfs(int u) {
    len[u] = 1; 
    for (int i = head[u]; i; i = e[i].nxt) {
        int v = e[i].to;
        if (v == fa[u]) continue;
        fa[v] = u;
        dfs(v);
        if (len[v] + 1 > len[u]) len[u] = len[v] + 1, son[u] = v;
    }
}
void solve(int u) {
    f[u][0] = 1;
    if (!son[u]) return ;
    int v = son[u];
    g[v] = g[u] - 1; f[v] = f[u] + 1;
    solve(v);
    ans += g[u][0];
    for (int i = head[u]; i; i = e[i].nxt) {
        v = e[i].to;
        if (v == fa[u] || v == son[u]) continue;
        g[v] = gp + len[v] + 1; gp += (len[v] * 2);
        f[v] = fp + 1; fp += len[v];
        solve(v);
        for (int j = 0; j <= len[v]; ++j) {
            if (j) ans += g[u][j] * f[v][j - 1];
            ans += f[u][j] * g[v][j + 1];
            if (j) g[u][j] += f[u][j] * f[v][j - 1];
            g[u][j] += g[v][j + 1];
            if (j) f[u][j] += f[v][j - 1];
        }
    }
}
int main() {
    int n = read();
    for (int i = 1; i < n; ++i) {
        int u = read(), v = read();
        add_edge(u, v);
    }
    dfs(1);
    g[1] = gp + len[1] + 1; gp += (len[1] * 2);
    f[1] = fp + 1; fp += len[1];
    solve(1);
    cout << ans;
    return 0;
}