CF-915-F-并查集

915-F 题目大意

给定一棵\(n\)个节点的树，节点带权，设函数\(I(x,y)\)等于点\(x\)到点\(y\)的路径上最大的点权与最小的点权之差。求：

\[\sum_{i=1}^{n}\sum_{j=i}^{n}I(i,j) \]

---

Solution

令函数\(F(i,j)\)等于点\(x\)到点\(y\)的路径上最大的点权，函数\(G(i,j)\)等于点\(x\)到点\(y\)的路径上最小的点权。所求的式子可转换为：

\[\sum_{i=1}^{n}\sum_{j=i}^{n}F(i,j)-\sum_{i=1}^{n}\sum_{j=i}^{n}G(i,j) \]

先求所有点对的\(F(i,j)\)之和，按点权从小到大枚举点\(x\)，枚举\(x\)连接的连通块，这些连通块的大小两两相乘，从而能够得到点\(x\)作为最大值的贡献，枚举的同时用并查集合并这些边。\(G(i,j)\)之和的求法相同。

时间复杂度\(O(nlogn)\)，瓶颈在排序上。

#include<bits/stdc++.h>
using namespace std;
using ll=long long;

template<typename T>
struct DSU{
	int n;
	vector<T> p,siz;
	DSU(int _n):p(_n+1),siz(_n+1),n(_n){
		iota(p.begin(),p.end(),0);
		for(int i=0;i<=n;i++) siz[i]=1;
	}
	T findd(T x){
		return p[x]==x?x:p[x]=findd(p[x]);
	}
	void unionn(T x,T y){
		x=findd(x),y=findd(y);
		if(x==y) return;
		if(siz[x]>siz[y]) swap(x,y);
		p[x]=y;
		siz[y]+=siz[x];
	}
};

void solve(){
    int n;
    cin>>n;
    vector<int> a(n);
    for(int i=0;i<n;i++){
        cin>>a[i];
    }
    vector<vector<int>> e(n);
    for(int i=1;i<n;i++){
        int x,y;
        cin>>x>>y;
        x--,y--;
        e[x].push_back(y);
        e[y].push_back(x);
    }
    vector<int> ids(n),vis(n);
    iota(ids.begin(),ids.end(),0);
    sort(ids.begin(),ids.end(),[&](const auto &i,const auto &j){
        return a[i]<a[j];
    });
    ll ans=0;
    DSU<int> mx(n),mn(n);
    for(int i=0;i<n;i++){
        int x=ids[i];
        int cnt=1;
        for(auto y:e[x]){
            if(!vis[y]) continue;
            ans+=1LL*cnt*mx.siz[mx.findd(y)]*a[x];
            cnt+=mx.siz[mx.findd(y)];
            mx.unionn(x,y);
        }
        vis[x]=1;
    }
    vis.assign(n,0);
    for(int i=n-1;~i;i--){
        int x=ids[i];
        int cnt=1;
        for(auto y:e[x]){
            if(!vis[y]) continue;
            ans-=1LL*cnt*mn.siz[mn.findd(y)]*a[x];
            cnt+=mn.siz[mn.findd(y)];
            mn.unionn(x,y);
        }
        vis[x]=1;
    }
    cout<<ans;
}

int main(){
    ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
    //freopen("input.txt","r",stdin);
	//freopen("output.txt","w",stdout);
    int T=1;
    //cin>>T;
    while(T--){
        solve();
    }
    return 0;
}