CF915F Imbalance Value of a Tree 题解

题意：

给定一棵 $n$ 个点的树，每个点有点权 $a_i$ 。

定义 $I(x,y)$ 为 $x$ 到 $y$ 的简单路径上的 $\max\{a_i\}-\min\{a_i\}$

求
$\sum_{i=1}^{n}\sum_{j=1}^{n} I(i,j)$
数据范围：

$1 \le n \le 10^6$

这个 $I(i,j)$ 是真的难看。考虑化简。

$\sum_{i=1}^{n}\sum_{j=1}^{n} f(i,j)-\sum_{i=1}^{n}\sum_{j=1}^{n} g(i,j)$

$f(i,j)$ 表示 $i$ 到 $j$ 的简单路径上的 $\max\{a_i\}$

$g(i,j)$ 表示 $i$ 到 $j$ 的简单路径上的 $\min\{a_i\}$

套路地，由于我们只需要知道总答案，而并非每一条路径的答案

于是考虑每个点的贡献。但是这个点就很不好处理。

我们可以先看看最熟悉的边权怎么处理

不失一般性，考虑边权最大值的贡献。

如果是边权的话，我们可以把所有的边从小到大排序

然后一条条边的加入，计算贡献后合并边两端的结点

对于当前的边 $(u,v)$ ，它的贡献就是

$w_{\text{mx}}(u,v) \times \mathrm{size}(u) \times \mathrm{size}(v)$

其中 $\mathrm{size}(u)$ 表示当前 $u$ 所在连通块的大小。

然后我们考虑点权怎么处理

考虑一条路径 $a \rightarrow b \rightarrow c$

它的点权最大值就是（~~懒得写 $a_a,a_b,a_c$ ，看得懂就好~~）

$\max\{a,b,c\}$

根据 $\max$ 多次嵌套不变的性质，可知

$\max\{a,b,c\}=\max\left\{\max\{a,b\}, ~\max\{b,c\}\right\}$

注意到这构成了 $(a,b),(b,c)$ 的配对，于是我们就可以把点权化为边权了。

$w_{\text{mx}}(u,v) \leftarrow \max\{a_u,a_v\}$

最小值也是一样的，就是从大到小排序而已。

时间复杂度 $O(n)$

树上的题目，如果 $n\le 10^6$ ，一般都是 $O(n)$ 的

代码：

#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdarg>
#include <cmath>
#include <iomanip>
#include <random>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
#define N ((int)(1e6+15))

int n,res,f[N],sz[N],val[N];
struct Edge{int x,y,mx,mn;}e[N];
void add(int &x,int y) {x += y;}
void init(int n){for(int i=1; i<=n; i++) f[i]=i, sz[i]=1;}
int find(int x){return f[x] == x ? x : f[x] = find(f[x]);}
void merge(int u,int v)
{
    int x=find(u),y=find(v);
    if(sz[x] > sz[y]) swap(x,y);
    f[x]=y; sz[y] += sz[x];
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    cin >> n;
    for(int i=1; i<=n; i++) cin >> val[i];
    for(int i=1; i<n; i++)
    {
        cin >> e[i].x >> e[i].y;
        e[i].mx = max(val[e[i].x],val[e[i].y]);
        e[i].mn = min(val[e[i].x],val[e[i].y]);
    }
    init(n); sort(e+1,e+n,[](Edge a,Edge b){return a.mn>b.mn;});
    for(int i=1,x,y; i<n; i++)
    {
        x=find(e[i].x); y=find(e[i].y);
        add(res, -e[i].mn * sz[x] * sz[y]); merge(x,y);
    }
    init(n); sort(e+1,e+n,[](Edge a,Edge b){return a.mx<b.mx;});
    for(int i=1,x,y; i<n; i++)
    {
        x=find(e[i].x); y=find(e[i].y);
        add(res, e[i].mx * sz[x] * sz[y]); merge(x,y);
    }
    cout << res << '\n';
    return 0;
}