hihocoder #1112 树上的好路径

时间限制:1000ms
单点时限:1000ms
内存限制:256MB

描述

现在有一棵有N个带权顶点的树，顶点编号为1,2,...,N。我们定义一条路径的次小(最小)权为它经过的所有顶点（包括起点和终点）中权值次小(最小)顶点的权值。现在给定常数c，你需要求出：存在多少个使得u<v的顶点组(u,v)，满足从u到v的最短路的次小权恰为c但最小权不为c。
输入

第一行有两个数N和c。(1<=n<=100000)

第二行N个数，依次表示每个顶点的权值。

接下来N-1行，每行两个数，代表这棵树的一条边所连接的两个顶点的编号。

我们保证输入中的数都在int以内。
输出

一个数，为答案。
样例输入

    8 2
    2 2 3 3 1 2 3 2
    1 2
    3 2
    3 8
    4 2
    5 2
    5 6
    6 7

样例输出

    17

---

Solution

为了方便, 把我们要考虑的树记作$T=(V, E)$, 用$w[u]$表示节点$u$ ($uin V$) 的权值.

先考虑一个简化的问题:

求最小权小于$c$且次小权不小于$c$的路径$(u, v)$的数目.

为了解决这个问题, 我们考虑如下的添边过程:

我们考虑一个动态的图$S(V, E'), E'subseteq E$.

从$S=(V, emptyset)$开始, 先把所有满足$w[u]ge c land w[v] ge c$的边$(u, v)$加到$S$中,

然后考虑满足

[w[u]<c land w[v]ge c lor w[u]ge c land w[v] <c]

的边$(u, v)$, 不失一般性, 不妨设 $w[u]<c, w[v]ge c$.

我们先把$u$固定为$u_0$, 考虑若将所有符合上述条件的边${(u_0, v)}$加到$s$中将能获得多少满足条件的路径.

显然这些满足条件的路径上的最小权就是$w[u_0]$.

(未完待续...)

(无力写了, 先把代码贴上)

---

UPD

前面写得太罗嗦了, 结果现在自己都看不大懂了. 其实做法一句话就能说清楚:

最小权小于$c$, 次小权不小于$c$的路径数 $-$ 最小权小于$c$, 次小权大于$c$的路径数

Implementation

#include <bits/stdc++.h>
using namespace std;
using LL=long long;
const int N{1<<17};

int a[N];

struct edge{
    int u, v;
    void read(){
        cin>>u>>v;
    }
}e[N];

struct DSU{
    int par[N], size[N];
    int n;
    DSU(int n):n(n){}
    void init(){
        for(int i=1; i<=n; i++){
            par[i]=i;
            size[i]=1;
        }
    }
    int find(int x){
        return x==par[x]?x:par[x]=find(par[x]);
    }
    void unite(int x, int y){
        x=find(x), y=find(y);
        if(x!=y) par[x]=y, size[y]+=size[x];
    }
};

vector<int> f[N];

void prep(DSU &b, int n, int c){
    b.init();
    for(int i=1; i<=n; i++) f[i].clear();
    for(int i=1; i<n; i++){
        int u=e[i].u, v=e[i].v;
        if(a[u]>=c && a[v]>=c){
            b.unite(u, v);
        }
    }
}

int main(){
    int n, c;
    cin>>n>>c;
    DSU b(n);

    for(int i=1; i<=n; i++)
        cin>>a[i];
    for(int i=1; i<n; i++)
        e[i].read();


    LL res=0;

    prep(b, n, c);

    for(int i=1; i<n; i++){
        int u=e[i].u, v=e[i].v;
        if(a[u]<c ^ a[v]<c){    //tricky
            // cout<<u<<' '<<v<<endl;
            if(a[v]<c) swap(u, v);
            int rv=b.find(v);
            // res+=LL(b.size[u])*LL(b.size[v]);
            // if(ru!=rv)
            f[u].push_back(b.size[rv]);
        }
    }

    for(int i=1; i<=n; i++){
        // if(f[i].size()) cout<<"#"<<i<<endl;
        LL sum=0, t=0;
        for(auto &x: f[i])
            sum+=x;
        for(auto &x: f[i]) t+=LL(x)*(sum-x);
        res+=t>>1;
        res+=sum;
    }


    prep(b, n, c+1);

    for(int i=1; i<n; i++){
        int u=e[i].u, v=e[i].v;
        if(a[u]<c && a[v]>c || a[u]>c && a[v]<c){    //tricky
            if(a[v]<c) swap(u, v);
            int rv=b.find(v);
            // res+=LL(b.size[u])*LL(b.size[v]);
            f[u].push_back(b.size[rv]);
        }
    }

    for(int i=1; i<=n; i++){
        LL sum=0, t=0;
        for(auto &x: f[i])
            sum+=x;
        for(auto &x: f[i]) t+=LL(x)*(sum-x);
        res-=t>>1, res-=sum;
    }

    cout<<res<<endl;
}