제출 #1358314

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1358314		goulthen	Chase (CEOI17_chase)	C++20	컴파일 에러	0 ms	0 KiB

chase

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

#define int long long
#define rep(i,a,b) for (int i = a; i <= b; i++)
#define per(i,a,b) for (int i = a; i >= b; i--)
#define pii pair<int,int>
#define fi first
#define se second
#define pb push_back

const int MAXN = 1e5+10;
const int INF = 1e9+67;
const int MOD = 1e9+7;
int a[MAXN], val[MAXN],k,ans,dp[MAXN][200];
vector<int> g[MAXN];

struct Top2 { 
    int b1=0,b2=0;
    int i1=-1;
    void insert(int x, int i) {
        if(x>b1) b2 = b1, i1 = i, b1 = x;
        else if(x>b2) b2 = x;
    }
    int ex(int i) {
        if(i==i1) return b2;
        return b1;
    }
    void reset() {b1=0;b2=0;i1=-1;}
} t2[MAXN][200];

void calc(int u, int p = -1) {
    for(int &v : g[u]) if(v!=p) { val[u] += a[v]; calc(v,u); }
}

void recompute(int u) {
    rep(ki,1,k) dp[u][ki] = max({t2[u][ki].b1, t2[u][ki-1].b1+val[u]});
}

void dfs(int u,int p = -1) {
    for (int &v : g[u]) if (v!=p) {
        dfs(v,u);
        rep(ki,1,k) t2[u][ki].insert(dp[v][ki],v);
    }
    recompute(u);
}


void reroot(int u, int p = -1) {
    ans = max(ans,dp[u][k]);
    for (int &v : g[u]) if(v!=p) {
        val[v] += a[u];
        val[u] -= a[v];
        rep(ki,1,k) {
            dp[u][ki] = max(t2[u][ki].ex(v),t2[u][ki-1].ex(v)+val[u]);
        }
        rep(ki,1,k) dp[v][ki] = max({dp[v][ki], dp[u][ki],dp[u][ki-1]+val[v]});

        reroot(v,u);

        val[v]-=a[u];
        val[u]+=a[v];
        recompute(u);
        recompute(v);
    }
}

void solve() {
    int n; cin >> n >> k;
    rep(i,1,n) cin >> a[i];
    rep(i,1,n-1) {
        int u,v; cin >> u >> v;
        g[u].pb(v);
        g[v].pb(u);
    }
    calc(1);
    dfs(1);
    reroot(1);

    cout << ans << endl;
}

int32_t main() {
    ios_base::sync_with_stdio(0);cin.tie(nullptr);
    int tt=1; 
    //cin >> tt;

    while (tt--) solve();
}