Submission #561310

#	Time	Username	Problem	Language	Result	Execution time	Memory
561310		fatemetmhr	Chase (CEOI17_chase)	C++17	100 / 100	642 ms	346044 KiB

chase

#include <bits/stdc++.h>

using namespace std;

typedef long long    ll;
typedef pair<ll, ll> pll;

#define pb      push_back
#define all(x)  x.begin(), x.end()
#define fi      first
#define se      second
#define mp      make_pair

const int maxn5  = 1e5 + 5;
const int mxnoon = 105;
const ll  inf    = 1e18;


ll valbanoon[maxn5];
ll a[maxn5], limlim, ans = 0;
vector <int> adj[maxn5];
ll dp_az_oon[2][maxn5][mxnoon], dp_be_oon[2][maxn5][mxnoon]; 

// dp[1] -> akharesh noon hast, 0 -> nist
// dp barabare ba TOM - JERRY

inline void dfs_be_oon_ja(int v, int par){
    dp_be_oon[1][v][0] = -inf;
    for(int lim = 1; lim <= limlim; lim++){
        dp_be_oon[1][v][lim] = valbanoon[v] - (par == -1 ? 0 : a[par]) - a[v];
        dp_be_oon[0][v][lim] = a[v];
    }
    dp_be_oon[0][v][0] = a[v];
    for(auto u : adj[v]) if(u != par){
        dfs_be_oon_ja(u, v);
        for(int lim = 0; lim <= limlim; lim++){
            dp_be_oon[0][v][lim] = max(dp_be_oon[0][v][lim], dp_be_oon[1][u][lim] + a[v]); // farz shode ishoon avale araye nistan
            dp_be_oon[0][v][lim] = max(dp_be_oon[0][v][lim], dp_be_oon[0][u][lim] - a[u] + a[v]);
            if(lim)
                dp_be_oon[1][v][lim] = max(dp_be_oon[1][v][lim], dp_be_oon[1][u][lim - 1] + valbanoon[v] - (par == -1 ? 0 : a[par]) - a[v]);
            if(lim)
                dp_be_oon[1][v][lim] = max(dp_be_oon[1][v][lim], dp_be_oon[0][u][lim - 1] + valbanoon[v] - (par == -1 ? 0 : a[par]) - a[u] - a[v]);
        }
    }
    return;
}

inline void dfs_az_oon_ja(int v, int par){
    for(auto u : adj[v]) if(u != par){
        dfs_az_oon_ja(u, v);
    }
    fill(dp_az_oon[1][v], dp_az_oon[1][v] + limlim + 2, valbanoon[v] - a[v]);
    dp_az_oon[1][v][0] = -inf;
    memset(dp_az_oon[0][v], 0, sizeof dp_az_oon[0][v]); // inja sare arayast .. pas har do mibinn
    for(auto u : adj[v]) if(u != par){
        for(int lim = 0; lim <= limlim; lim++){ // lim rooye tedad nan ha dar ta root
            ans = max(ans, dp_az_oon[0][v][lim] + dp_be_oon[0][u][limlim - lim] - a[u]);
            ans = max(ans, dp_az_oon[0][v][lim] + dp_be_oon[1][u][limlim - lim]);
            ans = max(ans, dp_az_oon[1][v][lim] + dp_be_oon[0][u][limlim - lim] - a[u]);
            ans = max(ans, dp_az_oon[1][v][lim] + dp_be_oon[1][u][limlim - lim]);
        }
        for(int lim = 0; lim <= limlim; lim++){
            dp_az_oon[0][v][lim] = max(dp_az_oon[0][v][lim], dp_az_oon[1][u][lim]);
            dp_az_oon[0][v][lim] = max(dp_az_oon[0][v][lim], dp_az_oon[0][u][lim]);
            if(lim)
                dp_az_oon[1][v][lim] = max(dp_az_oon[1][v][lim], dp_az_oon[1][u][lim - 1] - a[v] + valbanoon[v] - a[u]);
            if(lim)
                dp_az_oon[1][v][lim] = max(dp_az_oon[1][v][lim], dp_az_oon[0][u][lim - 1] + valbanoon[v] - a[u] - a[v]);
        }
        for(int lim = 1; lim <= limlim; lim++){
            dp_az_oon[0][v][lim] = max(dp_az_oon[0][v][lim], dp_az_oon[0][v][lim - 1]);
            dp_az_oon[1][v][lim] = max(dp_az_oon[1][v][lim], dp_az_oon[1][v][lim - 1]);
        }
    }
    reverse(all(adj[v]));
    fill(dp_az_oon[1][v], dp_az_oon[1][v] + limlim + 2, valbanoon[v] - a[v]);
    dp_az_oon[1][v][0] = -inf;
    memset(dp_az_oon[0][v], 0, sizeof dp_az_oon[0][v]); // inja sare arayast .. pas har do mibinn
    for(auto u : adj[v]) if(u != par){
        for(int lim = 0; lim <= limlim; lim++){ // lim rooye tedad nan ha dar ta root
            ans = max(ans, dp_az_oon[0][v][lim] + dp_be_oon[0][u][limlim - lim] - a[u]);
            ans = max(ans, dp_az_oon[0][v][lim] + dp_be_oon[1][u][limlim - lim]);
            ans = max(ans, dp_az_oon[1][v][lim] + dp_be_oon[0][u][limlim - lim] - a[u]);
            ans = max(ans, dp_az_oon[1][v][lim] + dp_be_oon[1][u][limlim - lim]);
        }
        for(int lim = 0; lim <= limlim; lim++){
            dp_az_oon[0][v][lim] = max(dp_az_oon[0][v][lim], dp_az_oon[1][u][lim]);
            dp_az_oon[0][v][lim] = max(dp_az_oon[0][v][lim], dp_az_oon[0][u][lim]);
            if(lim)
                dp_az_oon[1][v][lim] = max(dp_az_oon[1][v][lim], dp_az_oon[1][u][lim - 1] - a[v] + valbanoon[v] - a[u]);
            if(lim)
                dp_az_oon[1][v][lim] = max(dp_az_oon[1][v][lim], dp_az_oon[0][u][lim - 1] + valbanoon[v] - a[u] - a[v]);
        }
        for(int lim = 1; lim <= limlim; lim++){
            dp_az_oon[0][v][lim] = max(dp_az_oon[0][v][lim], dp_az_oon[0][v][lim - 1]);
            dp_az_oon[1][v][lim] = max(dp_az_oon[1][v][lim], dp_az_oon[1][v][lim - 1]);
        }
    }
    ans = max({ans, dp_az_oon[0][v][limlim], dp_az_oon[1][v][limlim]});
    return;
}

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

    int n; cin >> n >> limlim;
    for(int i = 0; i < n; i++){
        cin >> a[i];
        valbanoon[i] = a[i];
    }
    for(int i = 0; i < n - 1; i++){
        int a, b; cin >> a >> b;
        a--; b--;
        adj[a].pb(b);
        adj[b].pb(a);
    }

    for(int i = 0; i < n; i++) for(auto u : adj[i])
        valbanoon[i] += a[u];

    dfs_be_oon_ja(0, -1);
    dfs_az_oon_ja(0, -1);

    if(limlim){
        for(int i = 0; i < n; i++)
            ans = max(ans, valbanoon[i] - a[i]);
    }

    cout << ans << endl;
}

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