洛谷P5904 [POI2014] HOT-Hotels 加强版题解

题目链接：P5904 [POI2014] HOT-Hotels 加强版

题意：

给出一棵有 $n$ 个点的树，求有多少组点 $(i,j,k)$ 满足 $i,j,k$ 两两之间的距离都相等。

$(i,j,k)$ 与 $(i,k,j)$ 算作同一组。

输入格式：

第一行一个整数 $n$。

接下来 $n-1$ 行，每行两个整数 $a,b$，表示在 $a,b$ 之间有一条边。

输出格式：

一行一个整数，表示所有合法的点的组数。

数据范围：

对于 $100\%$ 的数据， $1\le n\le10^5, 1\le a\le b\le n$。

设 $f(i,j)$ 为 $i$ 所在子树中 $\mathrm{dis}(x,i)=j$ 的 $x$ 的个数。

设 $g(i,j)$ 为 $i$ 所在子树中满足

$\mathrm{dis}(\mathrm{LCA}(x, y), x)= \mathrm{dis}(\mathrm{LCA}(x, y), y)= \mathrm{dis}(\mathrm{LCA}(x, y), i)+j$

的无序二元组 $(x,y)$ 的个数。并记 $\mathrm{Ans}$ 为答案，那么

$\begin{aligned} \operatorname{Ans} &\operatorname{\big\uparrow} g(i,0) + \sum_{x,y \in \mathrm{son}(i) \,\land\, x \ne y} f(x,j-1)\cdot g(y,j+1) \\[6pt] g(i,j) &\operatorname{\big\uparrow} \sum_{x \in \mathrm{son}(i)} g(x,j+1) + \sum_{x,y \in \mathrm{son}(i) \,\land\, x \ne y} f(x,j-1)\cdot f(y,j-1) \\[6pt] f(i,j) &\operatorname{\big\uparrow} \sum_{x \in \mathrm{son}(i)} f(x,j-1) \end{aligned}$

这里 $\uparrow$ 均表示增量。直接转移的话，时间复杂度为 $\mathcal{O}(n^2)$ 。

注意到这里的第二维为深度，而每个节点所在子树的深度其实很多时候并没有 $\mathcal{O}(n)$ 。

考虑利用长链剖分优化转移，并利用指针直接 $\mathcal{O}(1)$ 继承重儿子的信息。

因为每个节点只属于一条长链，所以长链的总长度为 $\mathcal{O}(n)$ 。

对于每条轻链只会向重链合并一次，这相当于删去了轻链，故均摊也是 $\mathcal{O}(n)$ 的。

故总时间复杂度为 $\mathcal{O}(n)$

代码：

#include <bits/stdc++.h>
using namespace std;
// #define int long long
// #define INF 0x3f3f3f3f3f3f3f3f
typedef long long ll;
void up(int &x, int y) { x < y ? x = y : 0; }
void down(int &x, int y) { x > y ? x = y : 0; }
#define rep(i, a, b) for(int i = (a), i##END = (b); i <= i##END; i++)
#define Rep(i, a, b) for(int i = (a), i##END = (b); i >= i##END; i--)
#define pb push_back
#define N ((int)(1e5 + 15))

vector<int> vec[N];
int son[N], mxdep[N]; ll res, *f[N], *g[N], p[N * 4], *o = p;
void dfs(int u, int fa)
{
    for(auto v : vec[u]) if(v != fa)
        if(dfs(v, u), mxdep[v] > mxdep[son[u]]) son[u] = v;
    mxdep[u] = mxdep[son[u]] + 1;
}
void DP(int u, int fa)
{
    if(son[u]) { f[son[u]] = f[u] + 1; g[son[u]] = g[u] - 1; DP(son[u], u); }
    f[u][0] = 1; res += g[u][0];
    for(auto v : vec[u]) if(v != fa && v != son[u])
    {
        f[v] = o; o += mxdep[v] * 2; g[v] = o; o += mxdep[v] * 2;
        DP(v, u);
        rep(i, 0, mxdep[v] - 1)
        {
            if(i) res += f[u][i - 1] * g[v][i];
            res += g[u][i + 1] * f[v][i];
        }
        rep(i, 0, mxdep[v] - 1)
        {
            g[u][i + 1] += f[u][i + 1] * f[v][i];
            if(i) g[u][i - 1] += g[v][i];
            f[u][i + 1] += f[v][i];
        }
    }
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    int n; cin >> n;
    for(int i = 1, u, v; i < n; i++)
    {
        cin >> u >> v;
        vec[u].pb(v); vec[v].pb(u);
    }
    dfs(1, 0);
    f[1] = o; o += mxdep[1] * 2;
    g[1] = o; o += mxdep[1] * 2;
    DP(1, 0); cout << res << '\n';
    return 0;
}