题解 | # 小 Q 与树#

一个很标准的淀粉质题。

题目所给的式子，我们可以看做一个树的重心，然后周围点与这个重心的贡献和，然后继续找这个重心的子树的重心，重复上述过程，就是淀粉质了。

但是我们需要注意一点，如果我们在计算贡献时，没有对所给式子的第一项进行从小到大的排序的话，时间复杂度是要退化为N^2的。关于贡献的求法，就是一个点所有子树的距离和，加上这些点的个数乘于该点到重心的距离就是贡献。

那么我们很愉快地解决了这道题，下次再见！！！

#include <bits/stdc++.h>

const int P = 998244353;

int main()
{
    std::cin.tie(0);
    std::ios::sync_with_stdio(false);
    int n; std::cin >> n;
    std::vector <int> val(n);
    std::vector <int> e[n];
    for(int i = 0; i < n; ++i) std::cin >> val[i];
    for(int i = 0; i < n - 1; ++i) {
        int u, v; std::cin >> u >> v;
        u--, v--;
        e[u].emplace_back(v);
        e[v].emplace_back(u);
    }
    int root = -1;
    long long ans = 0;
    std::vector <int> siz(n);
    std::vector <int> max_siz(n);
    std::vector <int> vis(n,0);
    std::vector <std::pair <int, int>> temp;
    std::function <void(int, int, int)> get_root = [&] (int u, int fa, int res) {
        siz[u] = 1; max_siz[u] = 0;
        for(auto v : e[u]) {
            if(v == fa || vis[v]) continue;
            get_root(v, u, res);
            siz[u] += siz[v];
            max_siz[u] = std::max(max_siz[u], siz[v]);
        }
        max_siz[u] = std::max(max_siz[u], res - siz[u]);
        if(root == -1 || max_siz[u] < max_siz[root]) root = u;
    };
    std::function <void(int, int, int)> get_dis = [&] (int u, int fa, int dep) {
        temp.push_back({val[u],dep});
        for(auto v : e[u]) {
            if(v == fa || vis[v]) continue;
            get_dis(v, u, dep + 1);
        }
    };
    auto cal = [&] (int u, int fa, int dep) {
        temp.clear();
        long long res =0 ;
        get_dis(u, fa, dep);
        std::sort(temp.begin(),temp.end());
        long long dis_sum = 0, dis_pre =0;
        for(auto v : temp) {
            dis_sum += v.second;
        }
        int len = (int) temp.size();
        for(int i = 0; i <len; ++i) {
            dis_pre += temp[i].second;
            res = (res + 1ll * temp[i].first * (dis_sum - dis_pre) % P) % P;
            res = (res + 1ll * temp[i].first * temp[i].second % P * (len - i - 1) % P) % P;
        }
        return 2ll * res % P;
    };
    std::function <void(int)> solve = [&] (int u) {
        ans += cal(u, -1, 0);
        ans %= P;
        vis[u] = 1;
        for(auto v : e[u]) {
            if(vis[v]) continue;
            ans -= cal(v, u, 1);
            ans += P;
            ans %= P;
            root = -1;
            get_root(v, u, siz[v]);
            solve(root);
        }
    };
    get_root(0, -1, n);
    solve(root);
    std::cout << ans << '\n';
    return 0;
}