牛客 - 血压游戏

传送门：牛客 - 血压游戏

思路分析

显然只有深度一样的才会影响，所以将深度相同的节点单独拿出来建立虚树进行dp

AC代码

#pragma GCC optimize(2)
#include <bits/stdc++.h>
#define fi first
#define se second
#define ll long long
#define pb push_back
#define mp make_pair
#define fun function
#define sz(x) (x).size()
#define lowbit(x) (x)&(-x)
#define all(x) (x).begin(),(x).end()
#define mem(a,b) memset(a,b,sizeof(a))

namespace FastIO {
#define BUF_SIZE 100000
#define OUT_SIZE 100000
    bool IOerror=0;
    inline char nc() {
        static char buf[BUF_SIZE],*p1=buf+BUF_SIZE,*pend=buf+BUF_SIZE;
        if(p1==pend) {
            p1=buf;
            pend=buf+fread(buf,1,BUF_SIZE,stdin);
            if(pend==p1) {
                IOerror=1;
                return -1;
            }
        }
        return *p1++;
    }
    inline bool blank(char ch) {
        return ch==' '||ch=='\n'||ch=='\r'||ch=='\t';
    }
    template<class T> inline bool read(T &x) {
        bool sign=0;
        char ch=nc();
        x=0;
        for(; blank(ch); ch=nc());
        if(IOerror)return false;
        if(ch=='-')sign=1,ch=nc();
        for(; ch>='0'&&ch<='9'; ch=nc())x=x*10+ch-'0';
        if(sign)x=-x;
        return true;
    }
    template<class T,class... U>bool read(T& h,U&... t) {
        return read(h)&&read(t...);
    }
#undef OUT_SIZE
#undef BUF_SIZE
};
using namespace std;
using namespace FastIO;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

const int INF = 0x3f3f3f3f;
const int N = 2e5+10;


vector<int>e[N],tmp[N];
int dep[N],f[N],top[N],son[N],siz[N],dfn[N],tim;

void dfs1(int u,int fa) {
    dep[u]=dep[fa]+1;
    tmp[dep[u]].pb(u);
    siz[u]=1;
    son[u]=0;
    f[u]=fa;
    for(auto v:e[u]) {
        if(v==fa) continue;
        dfs1(v,u);
        siz[u]+=siz[v];
        if(siz[v]>siz[son[u]]) son[u]=v;
    }
}

void dfs2(int u,int t) {
    top[u]=t;
    dfn[u]=++tim;
    if(!son[u]) return ;
    dfs2(son[u],t);
    for(auto v:e[u]) {
        if(v==f[u] || v==son[u]) continue;
        dfs2(v,v);
    }
}

int getlca(int x,int y) {
    while(top[x]!=top[y]) {
        if(dep[top[x]]<dep[top[y]]) swap(x,y);
        x=f[top[x]];
    }
    if(dep[x]>dep[y]) swap(x,y);
    return x;
}

vector<int>new_e[N];
int st[N],Top;

void add(int u,int v){
    new_e[u].pb(v);
}

void insert(int x) {
    if(Top <= 1) {
        st[++Top] = x;
        return ;
    }
    int lca = getlca(x, st[Top]);
    if(lca == st[Top]) {
        st[++Top] = x;
        return;
    }
    while(Top > 1 && dfn[st[Top - 1]] >= dfn[lca]) {
        add(st[Top - 1], st[Top]);
        --Top;
    }
    if(lca != st[Top]) add(lca, st[Top]), st[Top] = lca;
    st[++Top] = x;
}

ll a[N]; 
int n,s;

ll dfs(int u){
    if(sz(new_e[u])==0) return a[u];
    ll ans=0;
    for(auto v:new_e[u]){
        ll sum=dfs(v);
        if(sum){
            ans+=max(1LL,sum-(dep[v]-dep[u]));
        }
    }
    return ans;
}

void del(int u){
    for(auto v:new_e[u]) del(v);
    new_e[u].clear();
}

ll solve(int pos){
    if(pos>1) tmp[pos].pb(s);
    sort(all(tmp[pos]),[](int x,int y){
        return dfn[x]<dfn[y];
    });
    Top=0;
    for(auto i:tmp[pos]) insert(i);
    while(Top>1){
        add(st[Top-1],st[Top]);
        --Top;
    }
    
    ll ans=dfs(s);
    del(s);
    return ans>1?ans-1:ans;
}


signed main() {

#ifdef xiaofan
    freopen("1.in","r",stdin);
    freopen("1.out","w",stdout);
#endif
    
    
    read(n,s);
    for(int i=1;i<=n;i++) read(a[i]);
    for(int i=1;i<n;i++){
        int u,v;
        read(u,v);
        e[u].pb(v);
        e[v].pb(u);
    }
    dfs1(s,0);
    dfs2(s,s);
    int maxdep=0;
    for(int i=1;i<=n;i++) maxdep=max(maxdep,dep[i]);
    ll ans=0;
    for(int i=1;i<=maxdep;i++) ans+=solve(i);
    printf("%lld\n",ans);



    return 0;
}