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$ ($u\in 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{<<};

 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=; i<=n; i++){
             par[i]=i;
             size[i]=;
         }
     }
     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=; i<=n; i++) f[i].clear();
     for(int i=; 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=; i<=n; i++)
         cin>>a[i];
     for(int i=; i<n; i++)
         e[i].read();

     LL res=;

     prep(b, n, c);

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

     prep(b, n, c+);

     for(int i=; 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=; i<=n; i++){
         LL sum=, t=;
         for(auto &x: f[i])
             sum+=x;
         for(auto &x: f[i]) t+=LL(x)*(sum-x);
         res-=t>>, res-=sum;
     }

     cout<<res<<endl;
 }