Submission #614112

#	Time	Username	Problem	Language	Result	Execution time	Memory
614112		leaked	Chase (CEOI17_chase)	C++17	100 / 100	1332 ms	341036 KiB

chase

#include <bits/stdc++.h>

#define f first
#define s second
#define vec vector
#define m_p make_pair
#define pb push_back
#define all(x) (x).begin(),(x).end()
#define rall(x) (x).rbegin(),(x).rend()
#define pw(x) (1LL<<(x))
#define sz(x) (int)(x).size()
#define fast_ioi ios_base::sync_with_stdio(false);cin.tie(0);cout.tie(0);

using namespace std;
template<class T> bool umin(T &a,const T &b){return (a>b?a=b,1:0);}
template<class T> bool umax(T &a,const T &b){return (a<b?a=b,1:0);}
typedef long long ll;
typedef pair<int,int> pii;
const ll inf=1e18;
const int N=1e5+1;
//#define int long lon

int a[N],cnt,n;
vec<int> g[N];
ll sm[N];

ll ans=0;

ll dpu[N][102][2];
ll dpd[N][102][2];
/// i'ma on and how many

//
//void dfs(int v,int p){
//    for(auto &z : g[v]){
//        if(z==p) continue;
//        for(int t=0;t<2;t++){
//            for(int j=0;j<=cnt;j++){
//                /// skip
//                umax(dp[z][j][1],dp[v][j][t]);
//                if(j+1<=cnt){
//                    /// add z
//                    ll cost=dp[v][j][t]+sm[z]-a[v]-a[z];
//                    umax(dp[z][j+1][0],cost);
//                }
//            }
//        }
//        dfs(z,v);
//    }
//}


void dfs1(int v,int p){

    for(int j=1;j<=cnt;j++){
        for(int t=0;t<2;t++){
            dpd[v][j][t]=(t==0?sm[v]-a[v]:0);
            dpu[v][j][t]=(t==0?sm[v]-a[v]:0);
        }
    }

    dpu[v][0][1]=0;dpd[v][0][1]=0;

    for(auto &z : g[v]){
        if(z==p) continue;
        dfs1(z,v);
    }

//    cout<<"REC " <<v<<endl;
    for(auto &z : g[v]){
        if(z==p) continue;


        vec<vec<ll>> ndpu(cnt+1,vec<ll>(2,-inf));
        vec<vec<ll>> ndpd(cnt+1,vec<ll>(2,-inf));



        for(int t=0;t<2;t++){
            for(int j=0;j<=cnt;j++){
                /// 0, v v will be
                if(j+1<=cnt){
                    umax(ndpu[j+1][t],ndpu[j][t]);
                    umax(ndpu[j+1][0],dpu[z][j][t]-a[v]-a[z]+sm[v]);
                }
                /// 1
                umax(ndpu[j][1],dpu[z][j][t]);
            }
        }

        for(int t=0;t<2;t++){
            for(int j=0;j<=cnt;j++){
                /// 0, v v will be
                int valt=-(t==0? a[v] : 0);
                if(j+1<=cnt){
                    umax(ndpd[j+1][t],ndpd[j][t]);
                    umax(ndpd[j+1][0],dpd[z][j][t]-a[v]+sm[v]+valt);
                }
                /// 1
                umax(ndpd[j][1],dpd[z][j][t]+valt);
            }
        }

        for(int j=2;j<=cnt-1;j++){
            int k=cnt-j;
            umax(ans,dpd[v][j][0]+ndpu[k+1][0]-sm[v]+a[v]);
            umax(ans,dpu[v][j][0]+ndpd[k+1][0]-sm[v]+a[v]);
        }

//        cout<<"ZI "<<z<<endl;
        for(int j=1;j<=cnt-1;j++){
            umax(ans,ndpu[j][1]+dpd[v][cnt-j][1]);
            umax(ans,ndpd[j][1]+dpu[v][cnt-j][1]);
        }



        for(int j=0;j<=cnt;j++){
            for(int t=0;t<2;t++)
                umax(dpu[v][j][t],ndpu[j][t]),umax(dpd[v][j][t],ndpd[j][t]);
        }

    }

    umax(ans,dpd[v][cnt][1]);
    umax(ans,dpu[v][cnt][1]);
    umax(ans,dpu[v][cnt][0]);
    umax(ans,dpd[v][cnt][0]);


}

signed main(){
    fast_ioi;

    cin>>n>>cnt;

    for(int i=0;i<n;i++) cin>>a[i];

    for(int i=1;i<n;i++){
        int v,u;
        cin>>v>>u;--v;--u;
        g[v].pb(u);g[u].pb(v);
    }

    for(int i=0;i<n;i++){
        sm[i]+=a[i];
        for(auto &z : g[i])
            sm[i]+=a[z];
    }
//    cout<<sm[5]<<endl;
    auto upd=[&](){
        for(int v=0;v<n;v++){
            for(int h=0;h<=cnt;h++){
                for(int t=0;t<2;t++)
                    dpd[v][h][t]=dpu[v][h][t]=-inf;
            }
        }
    };
    upd();

    dfs1(0,0);


    cout<<ans;
    return 0;
}
/*
10 3
10 1 1 1 1 10 1 10 1 10
1 2
2 3
3 4
4 5
5 6
3 7
7 8
9 4
9 10

*/

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