Submission #119400

#	Time	Username	Problem	Language	Result	Execution time	Memory
119400		zubec	Chase (CEOI17_chase)	C++14	100 / 100	440 ms	265592 KiB

chase

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;

const int N = 100100;

ll dp[N][110], dp2[N][110], a[N];

vector <int> g[N];

int n, k;

ll ans = 0;

void dfs(int v, int p){

    ll mx1[101] = {0}, mx2[101] = {0};

    ll sum = 0;
    for (int i = 0; i < g[v].size(); i++){
        sum += a[g[v][i]];
    }
    dp[v][0] = dp2[v][0] = 0;
    dp2[v][1] = sum;
    for (int i = 0; i < g[v].size(); i++){
        int to = g[v][i];
        if (to == p)
            continue;
        dfs(to, v);
        for (int j = 1; j <= k; j++){
            dp[v][j] = max(dp[v][j], max(dp[to][j], dp[to][j-1] + sum - a[p]));
            dp2[v][j] = max(dp2[v][j], max(dp2[to][j], dp2[to][j-1] + sum - a[to]));
            if (dp[to][j] >= mx1[j]){
                mx2[j] = mx1[j];
                mx1[j] = dp[to][j];
            } else {
                mx2[j] = max(mx2[j], dp[to][j]);
            }
        }
    }
    for (int i = 0; i < g[v].size(); i++){
        int to = g[v][i];
        if (to == p)
            continue;
        ans = max(ans, sum + dp[to][k-1]);
        for (int j = 0; j <= k; j++){
            ll mx = max(dp2[to][j], (j == 0 ? 0 : dp2[to][j-1] + sum - a[to]));
            if (mx1[k-j] == dp[to][k-j]){
                ans = max(ans, mx + mx2[k-j]);
            } else {
                ans = max(ans, mx + mx1[k-j]);
            }
        }
    }
}


int main(){
    ios_base::sync_with_stdio(0);cin.tie(0);

    cin >> n >> k;
    for (int i = 1; i <= n; i++){
        cin >> a[i];
    }
    for (int i = 2; i <= n; i++){
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    if (k == 0)
        return cout << 0, 0;



    dfs(1, 0);

    cout << ans;

}

/**



*/

Compilation message (stderr)

chase.cpp: In function 'void dfs(int, int)':
chase.cpp:21:23: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for (int i = 0; i < g[v].size(); i++){
                     ~~^~~~~~~~~~~~~
chase.cpp:26:23: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for (int i = 0; i < g[v].size(); i++){
                     ~~^~~~~~~~~~~~~
chase.cpp:42:23: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for (int i = 0; i < g[v].size(); i++){
                     ~~^~~~~~~~~~~~~

• Subtask 1: 20 / 20
• Subtask 2: 20 / 20
• Subtask 3: 30 / 30
• Subtask 4: 30 / 30