Submission #965961

#	Time	Username	Problem	Language	Result	Execution time	Memory
965961		zaa	Travelling Trader (CCO23_day2problem2)	C++14	25 / 25	244 ms	82004 KiB

Main

#include<bits/stdc++.h>
#define ll long long
#define ls(x) ((x)*2)
#define rs(x) ((x)*2+1)
#define Debug(...) fprintf(stderr, __VA_ARGS__)
#define For(i,a,b) for(int i=a,i##end=b;i<=i##end;i++)
#define Rof(i,a,b) for(int i=a,i##end=b;i>=i##end;i--)
#define rep(i,  b) for(int i=1,i##end=b;i<=i##end;i++)
using namespace std;
const int N=2e5+5,Mod=998244353;
//char buf[(1<<21)+5],*p1,*p2;
//#define getchar() (p1==p2&&(p2=(p1=buf)+fread(buf,1,1<<21,stdin),p1==p2)?EOF:*p1++)
inline void chmx(int &x,int y){(x<y)&&(x=y);}
inline void chmn(int &x,int y){(x>y)&&(x=y);}
//typedef __int128_t i128;
//i128 _base=1;
//inline int mol(int x){return x-Mod*(_base*x>>64);}
//inline void Add(int &x,int y){x=mol(x+y+Mod);}
inline int read(){
	int f=0,x=0;
	char ch=getchar();
	while(!isdigit(ch)){f|=(ch=='-');ch=getchar();}
	while(isdigit(ch)){x=(x<<3)+(x<<1)+(ch^48);ch=getchar();}
	return f?-x:x;
}int n,k;
vector<int>q[N];
ll lian[N];
int fa[N],dep[N];
ll dis[N],sum[N];
int a[N];
ll dp[N][5]; 
vector<int>q2[N],q3[N];
void dfs2(int u,int faz){
	sum[u]=a[u];
	for(auto v:q[u]){
		if(v==faz) continue;
		q2[u].emplace_back(v);
		dfs2(v,u);
		sum[u]+=a[v];
	}
	For(t,0,3){
		sort(q2[u].begin(),q2[u].end(),[&](int i,int j){return dp[i][t]-a[i]>dp[j][t]-a[j];});
		For(i,0,min((int)q2[u].size(),5)-1) q3[u].emplace_back(q2[u][i]);
	}
	sort(q2[u].begin(),q2[u].end(),[&](int i,int j){return dp[i][2]>dp[j][2];});
	For(i,0,min((int)q2[u].size(),5)-1) q3[u].emplace_back(q2[u][i]);
	sort(q3[u].begin(),q3[u].end());
	q3[u].resize(unique(q3[u].begin(),q3[u].end())-q3[u].begin());
	For(t,0,3)dp[u][t]=a[u];
	for(auto v:q3[u]){
		dp[u][0]=max(dp[u][0],dp[v][2]+a[u]);
		for(int w:q3[u]){
			if(v==w) continue;
			dp[u][0]=max(dp[u][0],dp[v][1]+dp[w][0]+sum[u]-a[v]-a[w]);
		}
		dp[u][0]=max(dp[u][0],dp[v][0]+sum[u]-a[v]);
		dp[u][1]=max(dp[u][1],dp[v][3]+sum[u]-a[v]);
		dp[u][3]=max(dp[u][3],dp[v][1]+sum[u]-a[v]);
		for(auto w:q3[u]){
			if(v==w) continue;
			dp[u][2]=max(dp[u][2],dp[v][3]+dp[w][2]+sum[u]-a[v]-a[w]);
			for(auto t:q3[u]){
				if(t==w||t==v) continue;
				dp[u][2]=max(dp[u][2],dp[v][3]+dp[w][1]+dp[t][0]+sum[u]-a[v]-a[w]-a[t]);
			}
		}
	}
	dp[u][2]=max(dp[u][2],dp[u][0]);
}
int tot,ans[N];
void geta(int u,int k){
	if(k==0||(k==2&&dp[u][2]==dp[u][0])){
		ans[++tot]=u;
		for(auto v:q3[u]){
			if(dp[u][0]==dp[v][0]+sum[u]-a[v]){
				for(auto w:q2[u]) if(v!=w) ans[++tot]=w;
				geta(v,0);
				return;
			}
			if(dp[u][0]==dp[v][2]+a[u]){
				geta(v,2);
				return;
			}
		}
		for(auto v:q3[u])for(auto w:q3[u])
			if(v!=w&&dp[u][0]==dp[v][1]+dp[w][0]+sum[u]-a[v]-a[w]){
				geta(v,1);
				for(auto x:q2[u]) if(v!=x&&w!=x) ans[++tot]=x;
				geta(w,0);
				return;
			}
//		dp[u][0]=max(dp[u][0],dp[v][2]+a[u]);
//		for(int w:q3[u]){
//			if(v==w) continue;
//			dp[u][0]=max(dp[u][0],dp[v][1]+dp[w][0]+sum[u]-a[v]-a[w]);
//		}
//		dp[u][0]=max(dp[u][0],dp[v][0]+sum[u]-a[v]); 
		return;
	}
	if(k==1){
		for(auto v:q3[u]){
			if(dp[u][1]==dp[v][3]+sum[u]-a[v]){
				for(auto w:q2[u]) if(v!=w) ans[++tot]=w;
				geta(v,3);
				ans[++tot]=u; 
				return;
			} 
		}
		ans[++tot]=u; 
//		dp[u][1]=max(dp[u][1],dp[v][3]+sum[u]-a[v]);
		return;
	}
	if(k==3){
		ans[++tot]=u; 
		for(auto v:q3[u]){
			if(dp[u][3]==dp[v][1]+sum[u]-a[v]){
				geta(v,1);
				for(auto w:q2[u]) if(v!=w) ans[++tot]=w;
				return;
			} 
		}
//		dp[u][3]=max(dp[u][3],dp[v][1]+sum[u]-a[v]);
		return;
	}
	for(auto v:q3[u]){
		for(auto w:q3[u]){
			if(v==w) continue;
			if(dp[u][2]==dp[v][3]+dp[w][2]+sum[u]-a[v]-a[w]){
				for(auto s:q2[u]) if(v!=s&&s!=w) ans[++tot]=s;
				geta(v,3);
				ans[++tot]=u;
				geta(w,2);
				return;
			}
			for(auto t:q3[u]){
				if(t==w||t==v) continue;
				if(dp[u][2]==dp[v][3]+dp[w][1]+dp[t][0]+sum[u]-a[v]-a[w]-a[t]){
					geta(v,3);
					ans[++tot]=u;
					geta(w,1);
					for(auto s:q2[u]) if(v!=s&&t!=s&&w!=s) ans[++tot]=s;
					geta(t,0);
					return;
				}
			}
		}
	}
	
//	dp[u][2]=max(dp[u][2],dp[u][0]);
//	dp[u][2]=max(dp[u][2],dp[v][3]+dp[w][2]+sum[u]-a[v]-a[w]);
//	dp[u][2]=max(dp[u][2],dp[v][3]+dp[w][1]+dp[t][0]+sum[u]-a[v]-a[w]-a[t]);
	
}
void dfs(int x,int faz){
	fa[x]=faz;
	dep[x]=dep[faz]+1;
	lian[x]=a[x];
	dis[x]=dis[faz]+a[x];
	for(auto y:q[x]){
		if(y==faz) continue;
		dfs(y,x);
		lian[x]=max(lian[x],lian[y]+a[x]);
	}
}
void dfs1(int x,ll v){
	ans[++tot]=x;
	for(auto y:q[x])
		if(y!=fa[x]&&lian[y]==lian[1]-v-a[x])
			dfs1(y,v+a[x]);
}
bool vis[N];
void dfs3(int x,int v){
	if(v==0)vis[ans[++tot]=x]=1;
	int p=q[x].size()-1;
	for(auto y:q[x]){
		if(y==fa[x])continue;
		p--;
		if(p==(int)q[x].size()-2)dfs3(y,(v+1)%3);
		else{
			if(dep[ans[tot]]-dep[x]==1) dfs3(y,2);
			else dfs3(y,0);
		}
	}
	if(!vis[x])vis[ans[++tot]=x]=1;
}
inline void solve(){
	n=read(),k=read();
	For(i,1,n-1){
		int u=read(),v=read();
		q[u].push_back(v);
		q[v].push_back(u);
	}
	ll sum=0;
	For(i,1,n) a[i]=read(),sum+=a[i];
	if(k==1){
	dfs(1,0);
		dfs1(1,0);
		printf("%lld\n",lian[1]);
		printf("%d\n",tot);
		For(i,1,tot) printf("%d ",ans[i]);
		return ;
	}
	if(k==3){
		dfs(1,0);
		dfs3(1,0);
		printf("%lld\n",sum);
		printf("%d\n",tot);
		For(i,1,tot) printf("%d ",ans[i]);
		return ;
	}
	if(k==2){
		dfs2(1,0);
		tot=0;
		geta(1,0);
		printf("%lld\n",dp[1][0]);
		printf("%d\n",tot);
		For(i,1,tot) printf("%d ",ans[i]);
		return ;
	}
}
signed main(){
	//_base=(_base<<64)/Mod;
	//freopen(".in","r",stdin);
	//freopen(".out","w",stdout);
	// ios::sync_with_stdio(false);
	// cin.tie(0); cout.tie(0);
	int T=1;
	while(T--){solve();}
#ifdef LOCAL
    Debug("\nMy Time: %.3lfms\n", (double)clock() / CLOCKS_PER_SEC);
#endif
	return 0;
}

• Subtask 1: 2 / 2
• Subtask 2: 7 / 7
• Subtask 3: 3 / 3
• Subtask 4: 4 / 4
• Subtask 5: 4 / 4
• Subtask 6: 5 / 5