【BZOJ4543】Hotel加强版

【BZOJ4543】Hotel加强版

题面

bzoj

洛谷

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

题解

我们分析一下这个问题，要怎么样的点才满足三点距离两两相等呢？

1、存在三个点有共同的$LCA$。

2、存在一个点，使得它到它两颗不同的子树种两点的距离为$d$且它存在$d$级祖先。

考虑$dp:$

设$f[i][j]$表示以$i$为根的子树中，距离$i$为$j$的点数

$g[i][j]$表示以$i$为根的子树中两点到$LCA$距离为$d$,并且它们$LCA$到$j$的距离为$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$