Submission #164412

#	Time	Username	Problem	Language	Result	Execution time	Memory
164412		Lawliet	Chase (CEOI17_chase)	C++14	100 / 100	1683 ms	363112 KiB

chase

#include <bits/stdc++.h>

using namespace std;
typedef long long int lli;

const int MAXK = 110;
const int MAXN = 100010;

int n, k;

int v[MAXN];
int maxDP[MAXN][MAXK][2];

lli ans;

lli sum[MAXN];
lli dp[MAXN][MAXK];

vector< int > adj[MAXN];

void DFSCalculate(int cur, int p)
{
	for(int i = 0 ; i < adj[cur].size() ; i++)
	{
		int viz = adj[cur][ i ];

		if( viz == p ) continue;

		sum[ cur ] += v[ viz ];
		DFSCalculate( viz , cur );
	}

	for(int j = 1 ; j <= k ; j++)
	{
		for(int i = 0 ; i < adj[cur].size() ; i++)
		{
			int viz = adj[cur][ i ];

			if( viz == p ) continue;

			if( dp[viz][j] > dp[ maxDP[cur][j][1] ][ j ] ) maxDP[ cur ][ j ][1] = viz;
			if( dp[ maxDP[cur][j][1] ][ j ] > dp[ maxDP[cur][j][0] ][ j ]) 
				swap( maxDP[ cur ][ j ][0] , maxDP[ cur ][ j ][1] );
		}

		dp[ cur ][ j ] = dp[ maxDP[cur][j - 1][0] ][ j - 1 ] + sum[ cur ];
		dp[ cur ][ j ] = max( dp[ cur ][ j ] , dp[ maxDP[cur][j][0] ][ j ] );
	}
}

void DFSRotate(int cur, int p)
{
	ans = max( ans , dp[cur][k] );

	for(int i = 0 ; i < adj[cur].size() ; i++)
	{
		int viz = adj[cur][ i ];

		if( viz == p ) continue;

		vector< lli > oldDPCur;
		vector< lli > oldDPViz;

		vector< int > oldMaxDPCur[2];
		vector< int > oldMaxDPViz[2];

		lli oldSumCur = sum[ cur ];
		lli oldSumViz = sum[ viz ];

		for(int j = 0 ; j <= k ; j++)
		{
			oldDPCur.push_back( dp[cur][j] );
			oldDPViz.push_back( dp[viz][j] );

			for(int p = 0 ; p < 2 ; p++)
			{
				oldMaxDPCur[p].push_back( maxDP[cur][j][p] );
				oldMaxDPViz[p].push_back( maxDP[viz][j][p] );
			}
		}

		sum[ cur ] -= v[ viz ];
		sum[ viz ] += v[ cur ];

		for(int j = 1 ; j <= k ; j++)
		{
			lli newDPCur = 0;

			if( maxDP[ cur ][ j - 1 ][0] == viz ) newDPCur = dp[ maxDP[cur][j - 1][1] ][ j - 1 ] + sum[ cur ];
			else newDPCur = dp[ maxDP[cur][j - 1][0] ][ j - 1 ] + sum[ cur ];

			if( maxDP[ cur ][ j ][0] == viz ) newDPCur = max( newDPCur , dp[ maxDP[cur][j][1] ][ j ] );
			else newDPCur = max( newDPCur , dp[ maxDP[cur][j][0] ][ j ] );

			dp[ cur ][ j ] = newDPCur;
		}

		for(int j = 1 ; j <= k ; j++)
		{
			if( dp[cur][j] > dp[ maxDP[viz][j][1] ][ j ] ) maxDP[ viz ][ j ][1] = cur;
			if( dp[ maxDP[viz][j][1] ][ j ] > dp[ maxDP[viz][j][0] ][ j ]) 
				swap( maxDP[ viz ][ j ][0] , maxDP[ viz ][ j ][1] );

			dp[ viz ][ j ] = max( dp[ viz ][ j ] , dp[ maxDP[viz][j][0] ][ j ] );
			dp[ viz ][ j ] = max( dp[ viz ][ j ] , dp[ maxDP[viz][j - 1][0] ][ j - 1 ] + sum[ viz ]);
		}

		DFSRotate( viz , cur );

		sum[ cur ] = oldSumCur;
		sum[ viz ] = oldSumViz;

		for(int j = 0 ; j <= k ; j++)
		{
			dp[cur][j] = oldDPCur[j];
			dp[viz][j] = oldDPViz[j];

			for(int p = 0 ; p < 2 ; p++)
			{
				maxDP[cur][j][p] = oldMaxDPCur[p][j];
				maxDP[viz][j][p] = oldMaxDPViz[p][j];
			}
		}
	}
}

int main()
{
	scanf("%d %d",&n,&k);

	for(int i = 1 ; i <= n ; i++)
		scanf("%d",&v[i]);

	for(int i = 1 ; i < n ; i++)
	{
		int U, V;
		scanf("%d %d",&U,&V);

		adj[ U ].push_back( V );
		adj[ V ].push_back( U );
	}

	DFSCalculate( 1 , 0 );
	DFSRotate( 1 , 0 );

	printf("%lld\n",ans);
}

Compilation message (stderr)

chase.cpp: In function 'void DFSCalculate(int, int)':
chase.cpp:23:20: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for(int i = 0 ; i < adj[cur].size() ; i++)
                  ~~^~~~~~~~~~~~~~~~~
chase.cpp:35:21: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for(int i = 0 ; i < adj[cur].size() ; i++)
                   ~~^~~~~~~~~~~~~~~~~
chase.cpp: In function 'void DFSRotate(int, int)':
chase.cpp:55:20: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for(int i = 0 ; i < adj[cur].size() ; i++)
                  ~~^~~~~~~~~~~~~~~~~
chase.cpp: In function 'int main()':
chase.cpp:129:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  scanf("%d %d",&n,&k);
  ~~~~~^~~~~~~~~~~~~~~
chase.cpp:132:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d",&v[i]);
   ~~~~~^~~~~~~~~~~~
chase.cpp:137:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d %d",&U,&V);
   ~~~~~^~~~~~~~~~~~~~~

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