Submission #1358250

#TimeUsernameProblemLanguageResultExecution timeMemory
1358250goulthenChase (CEOI17_chase)C++20
40 / 100
4098 ms170088 KiB
#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], dp[MAXN][200], val[MAXN],k,ans;
vector<int> g[MAXN];

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

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

    rep(ki,1,k) {
        for(int &v : g[u]) dp[u][ki] = max({dp[u][ki], dp[v][ki], dp[v][ki-1] + val[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] = dp[v][ki] = 0;
            for(int &nv : g[u]) if(nv!=v) dp[u][ki] = max({dp[u][ki],dp[nv][ki],dp[nv][ki-1]+val[u]});
            for(int &nv : g[v]) dp[v][ki] = max({dp[v][ki], dp[nv][ki],dp[nv][ki-1]+val[v]});
        }
        reroot(v,u);
        val[v]-=a[u];
        val[u]+=a[v];
        rep(ki,1,k) {
            dp[u][ki] = dp[v][ki] = 0;
            for(int &nv : g[u]) dp[u][ki] = max({dp[u][ki],dp[nv][ki],dp[nv][ki-1]+val[u]});
            for(int &nv : g[v]) if(nv!=u) dp[v][ki] = max({dp[v][ki], dp[nv][ki],dp[nv][ki-1]+val[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();
}
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...