$\operatorname{牛半仙的妹子树}$

题目链接：https://ac.nowcoder.com/acm/problem/212916

到主站看：https://blog.csdn.net/weixin_43346722/article/details/109272007

题目

牛半仙有 $n_{}$ 个妹子，她们所在的位置组成一棵树，相邻两个妹子的距离为 $1_{}$ 。
有 $m_{}$ 个妹子具有超能力，她们会影响到其他妹子。
离所有具有超能力的妹子的最短距离在 $[l,r]_{}$ 间的妹子会受到影响，会具有旪超能力。
这些具有能力的妹子共同形成了一个磁场。对于一个位置，一个具有超能力的妹子为其增加的磁场强度为妹子到这个位置的距离的平方，一个具有旪超能力的妹子为其增加的磁场强度为妹子到这个位置的距离。
现在牛半仙想知道一个位置的磁场强度有多大。
因为牛半仙对妹子们特别关心，所以他有 $k_{}$ 个询问。

输入

第一行 $5_{}$ 个正整数，代表 $n,m,k,l,r_{}$ 。
第 $2_{}$ 行到第 $n_{}$ 行每行 $2_{}$ 个正整数 $u_i,v_i$ ，代表树中存在边 $(u_i,v_i)$ 。
第 $n+1_{}$ 行有 $m_{}$ 个正整数，代表这些妹子有超能力。
第 $n+2_{}$ 行有 $k_{}$ 个正整数，代表询问的妹子。对于每个询问，输出一行，代表该询问的答案。

输出

对于每个询问，输出一行，代表该询问的答案。

样例输入

样例输出

样例解释

在这里插入图片描述

图中红点为超能力,绿点为旪超能力

$1$ 号点权值为 $1^2+3^2+1+3=14$
$6$ 号点权值为 $1^2+5^2+3+5=34$
$9$ 号点权值为 $0^2+4^2+2+2=20$
$11$ 号点权值为 $1^2+5^2+3+3=32$

数据范围

对于 $5\%$ 的数据： $n,m,k≤100$
对于 $20\%$ 的数据： $n,m,k\le2000$
对于另 $20\%$ 的数据： $u_i=i,v_i=i+1$
对于另 $20\%$ 的数据： $u_i=1$
对于 $100\%$ 的数据： $n,m,k\le5\times10^5$

思路

这道题是一道很烦的题目，用了 dp、换根等等操作。

首先，我们可以求出某一个点是否是超能力点或者旪超能力点。
我这里用 $sup[i]$ 和 $psup[i]$ 表示。
那超能力点很好搞，至于旪超能力点，我们可以通过bfs，从每一个超能力点出发，算出每一个点距离最近的超能力点的距离 $dis[i]$ 。
（要用bfs因为这样可以保证第一次找到的时候距离最短）
那如果这个距离在 $[l,r]$ 之间，那这个点就是旪超能力点。

接着，我们可以通过dfs的树形dp求出一一为根的时候的一些数值： $sum[i]$ 表示点 $i$ 与它的子树中有多少个点是超能力点 $psum[i]$ 表示点 $i$ 与它的子树中有多少个点是旪超能力点
（这两个很好搞，就直接等于它每个子树的个数之和，如果自己本身是就 $+1$ ） $ddis[i]$ 表示点 $i$ 与它的子树所有超能力点对答案的贡献（ $i$ 就是要到的位置） $pdis[i]$ 表示点 $i$ 与它的子树所有旪超能力点对答案的贡献

$pdis[i]$ 还好说，就直接每一个子树的 $pdis[i]$ 加每一个子树的旪超能力点数。（因为每一个旪超能力点距离要到的位置都加了 $1$ ，那合在一起就是 $1*psum[i]$ 即 $psum[i]$ ）
但是 $ddis[i]$ 怎么的每一个数都要平方，怎么办呢？
先搞点简单的， $dis[i]$ 表示点 $i$ 与它的子树所有超能力点到 $i$ 的距离和。
那求这个的方式就和 $pdis[i]$ 一样，只是把 $pdis[i]$ 换成 $dis[i]$ ， $psum[i]$ 换成 $sum[i]$ 。

那接着我们来看如何求 $ddis[i]$ 。
设原来一个在子树 $x$ 中的超能力点距离子树 $x$ 的根是 $y$ 的距离为 $z$ ，那么上移到 $i$ ，距离加 $1$ ，贡献就变成了 $(z+1)^2$ ，化简得到 $z^2+2z+1$ 。
枚举每一个子树的根 $y$ ，那就是所有 $ddis[y]+2*dis[y]+sum[y]$ 加起来。

但是这只是一个点为要到的位置，如果正常枚举每一个点，就会超时。
这个时候，我们就要用到换根。（就是枚举儿子，把它拎到上面做根，然后维护儿子的各项值值）

那我们来看如何维护上面的值，设原来的根是 $x$ ，现在要把根变成它的一个儿子 $i$ 。

$sum[]$ 好弄，就 $sum[i]=sum[x]$ 。 $psum[]$ 也是一样： $psum[i]=psum[x]$

维护答案的思路大致是这样的：
在这里插入图片描述

$pdis[]$ 还好，就 $pdis[i] = pdis[i]+((pdis[x] - (pdis[i] + psum[i])) + (psum[x] - psum[i]))$
（ $pdis[i]$ 就是蓝色部分的贡献）
（ $pdis[i]+psum[i]$ 表示新做根的点及下面的子树原来对答案的贡献）
（那 $pdis[x] - (pdis[i] + psum[i])$ 就是不是那棵子树的点原来对答案的贡献）
（ $psum[x] - psum[i]$ 就是不在蓝色框框的点的个数，也就是这些点走一个距离增加的贡献。因为之前只只用到 $x$ ，这一次要多走到 $i$ ，所以多走了一个距离）

那 $dis[]$ 同理，就只要改一下变量名即可。

接着我们来看最烦的 $ddis[]$ ，先把式子写出来： $diss[i]=ddis[i] + ((ddis[x] - (ddis[i] + 2 * dis[i] + sum[i])) + 2 * (dis[x] - (dis[i] + sum[i])) + (sum[x] - sum[i]))$
（ $ddis[i]$ 就是蓝色部分的贡献）
（ $ddis[i] + (2 * dis[i] + sum[i])$ 就是新做根的点及下面的子树原来对答案的贡献）
（那 $ddis[x] - (ddis[i] + 2 * dis[i] + sum[i])$ 就是不是那棵子树的点原来对答案的贡献）
（ $2 * (dis[x] - (dis[i] + sum[i])) + (sum[x] - sum[i])$ 就是不是蓝色部分的点走一个距离增加的贡献：由 $x^2$ 变成 $x^2+2x+1$ ，多了 $2x+1$ ）

比赛时

看到题目，想了半天都想不出比较优的解，就只能唯唯诺诺的打一个只有 $5$ 分的 Floyed。

图片说明

代码

#include<queue>
#include<cstdio>
#include<cstring>
#define ll long long

using namespace std;

struct node {
    ll to, nxt;
}e[5000001];
struct bfsu {
    ll x, size;
}noww;
ll n, m, k, l, r, x, y, le[500001], KK, sup[500001], psup[500001], in[500001], re;
ll sum[500001], psum[500001], dis[500001], pdis[500001], ddis[500001];
queue <bfsu> q;
char c;

ll read() {
    re = 0;
    c = getchar();
    while (c < '0' || c > '9') c = getchar();
    while (c >= '0' && c <= '9') {
        re = re * 10 + c - '0';
        c = getchar();
    }
    return re;
}

void add(ll x, ll y) {
    e[++KK] = (node){y, le[x]}; le[x] = KK;
    e[++KK] = (node){x, le[y]}; le[y] = KK;
}

void bfs() {//bfs求出每个点距离最近的超能力点的距离
    while (!q.empty()) {
        noww = q.front();
        q.pop();

        dis[noww.x] = noww.size;
        for (ll i = le[noww.x]; i; i = e[i].nxt)
            if (!in[e[i].to]) {
                in[e[i].to] = 1;
                q.push((bfsu){e[i].to, noww.size + 1});
            }
    }
}

void dfsget(ll now, ll fa) {//通过dp求出以1位根时的对应值
    sum[now] = sup[now];
    psum[now] = psup[now];

    for (ll i = le[now]; i; i = e[i].nxt)
        if (e[i].to != fa) {
            dfsget(e[i].to, now);

            sum[now] += sum[e[i].to];
            psum[now] += psum[e[i].to];

            ddis[now] += ddis[e[i].to] + 2 * dis[e[i].to] + sum[e[i].to];
            dis[now] += dis[e[i].to] + sum[e[i].to];
            pdis[now] += pdis[e[i].to] + psum[e[i].to];
        }
}

void dfschan(ll now, ll fa) {//换根
    for (ll i = le[now]; i; i = e[i].nxt)
        if (e[i].to != fa) {
            ddis[e[i].to] += (ddis[now] - (ddis[e[i].to] + 2 * dis[e[i].to] + sum[e[i].to])) + 2 * (dis[now] - (dis[e[i].to] + sum[e[i].to])) + (sum[now] - sum[e[i].to]);
            dis[e[i].to] += (dis[now] - (dis[e[i].to] + sum[e[i].to])) + (sum[now] - sum[e[i].to]);
            pdis[e[i].to] += (pdis[now] - (pdis[e[i].to] + psum[e[i].to])) + (psum[now] - psum[e[i].to]);
            sum[e[i].to] = sum[now];
            psum[e[i].to] = psum[now];

            dfschan(e[i].to, now);
        }
}

int main() {
    n = read();
    m = read();
    k = read();
    l = read();
    r = read();
    for (ll i = 1; i < n; i++) {
        x = read();
        y = read();
        add(x, y);
    }

    for (ll i = 1; i <= m; i++) {
        x = read();
        sup[x] = 1;
        q.push((bfsu){x, 0});
        in[x] = 1;
    }

    bfs();//bfs求出每个点距离最近的超能力点的距离
    for (ll i = 1; i <= n; i++)
        if (dis[i] >= l && dis[i] <= r)
            psup[i] = 1;//找到旪超能力点

    memset(dis, 0, sizeof(dis));
    dfsget(1, 0);//通过dp求出以1位根时的对应值
    dfschan(1, 0);//换根

    for (ll i = 1; i <= k; i++) {
        x = read();
        printf("%lld\n", ddis[x] + pdis[x]);
    }

    return 0;
}