Submission #161923

#	Time	Username	Problem	Language	Result	Execution time	Memory
161923		shayan_p	Chase (CEOI17_chase)	C++14	50 / 100	603 ms	260424 KiB

chase

// Remember...

#include<bits/stdc++.h>

#define F first
#define S second
#define PB push_back
#define sz(s) int((s).size())
#define bit(n,k) (((n)>>(k))&1)

using namespace std;

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

const int maxn=1e5+10, maxk=105;
const ll inf=1e18;

vector<int> v[maxn];
int a[maxn], n, k;
ll dn[maxn][maxk], up[maxn][maxk], arr[maxk], ANS;
int who[maxn][maxk][2];

void dfsdn(int u,int par=-1){
    memset(dn[u],0,sizeof dn[u]);
    ll num=0;
    for(int y:v[u]){
	if(y==par) continue;
	num+= a[y];
    }
    for(int y:v[u]){
	if(y==par) continue;
	dfsdn(y,u);
	for(int i=0;i<=k;i++)
	    dn[u][i]= max( dn[u][i], dn[y][i] );
	for(int i=0;i< k;i++)
	    dn[u][i+1]= max( dn[u][i+1], dn[y][i] + num );
    }
    for(int i=1;i<=k;i++)
	dn[u][i]= max(dn[u][i], num);
    for(int i=0;i<=k;i++)
	who[u][i][0]= who[u][i][1]= -1;
    for(int y:v[u]){
	if(y==par) continue;
	for(int i=0;i<=k;i++){
	    if(who[u][i][0]==-1 || dn[y][i]>dn[ who[u][i][0] ][i])
		who[u][i][1]= who[u][i][0], who[u][i][0]=y;
	    else if(who[u][i][1]==-1 || dn[y][i]>dn[ who[u][i][1] ][i])
		who[u][i][1]=y;
	}
    }
    ANS=max(ANS, dn[u][k]);
}
void dfs(int u,int par=-1){
    ll num=0;
    for(int y:v[u])
	num+= a[y];
    for(int y:v[u]){
	if(y==par) continue;
	dfs(y,u);
	for(int i=0;i<=k;i++){
	    arr[i]= up[y][i];
	}
	for(int i=1;i<=k;i++){
	    arr[i]= max(arr[i], up[y][i-1] + num - a[y]);
	}
	for(int i=0;i<=k;i++){
	    int vert= who[u][k-i][0] == y ? who[u][k-i][1] : who[u][k-i][0];
	    ANS=max(ANS, arr[i] + dn[vert][k-i]);
	}
	for(int i=0;i<=k;i++){
	    up[u][i]= max(up[u][i], arr[i]);
	}
    }
    for(int i=1;i<=k;i++){
	up[u][i]= max(up[u][i], num);
    }
    ANS=max(ANS, up[u][k]);
}

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

    cin>>n>>k;

    for(int i=1;i<=n;i++){
	cin>>a[i];
    }
    for(int i=0;i<n-1;i++){
	int a,b; cin>>a>>b;
	v[a].PB(b);
	v[b].PB(a);
    }
    dfsdn(1), dfs(1);
    return cout<<ANS<<endl,0;
}
// Deathly mistakes:
//  * Read the problem carefully.
//  * Check maxn.
//  * Overflows.


// #pragma GCC optimize("Ofast,no-stack-protector,unroll-loops,fast-math")
// #pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx,tune=native")

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