Submission #392022

#	Time	Username	Problem	Language	Result	Execution time	Memory
392022		AriaH	Chase (CEOI17_chase)	C++11	100 / 100	1504 ms	251440 KiB

chase

/** I can do this all day **/

#pragma GCC optimize("O2")
#include <bits/stdc++.h>
using namespace std;

typedef long long                   ll;
typedef long double                 ld;
typedef pair<int,int>               pii;
typedef pair<ll,ll>                 pll;
#define all(x)                      (x).begin(),(x).end()
#define F                           first
#define S                           second
#define Mp                          make_pair
#define SZ(x)			    		(int)x.size()
#define fast_io                     ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
#define file_io                     freopen("in.txt" , "r+" , stdin) ; freopen("out.txt" , "w+" , stdout);

const int N = 1e5 + 5;
const ll mod = 1e9 + 7;
const ll mod2 = 998244353;
const ll inf = 1e18;
const int LOG = 22;
const int M = 1e2 + 5;

int n, k, Par[N];

vector < int > G[N];

ll tot, A[N], nei[N], dp_down[2][M][N], dp_up[M][N];

void dfs_down(int v, int P)
{
	Par[v] = P;
	dp_up[1][v] = nei[v];
	for(auto u : G[v])
	{
		if(u == P) continue;
		dfs_down(u, v);
		for(int x = 0; x <= k; x ++)
		{
			dp_up[x][v] = max(dp_up[x][v], dp_up[x][u]);
			if(x) 
			{
				dp_up[x][v] = max(dp_up[x][v], dp_up[x - 1][u] + nei[v] - A[u]);
			}
		}
	}
}

void dfs(int v, int P)
{
	for(int i = 1; i <= k; i ++)
	{
		dp_down[1][i][v] = nei[v] - A[P];
	}
	dp_down[1][0][v] = -inf;
	for(auto u : G[v])
	{
		if(u == P) continue;
		dfs(u, v);
		for(int x = 0; x <= k; x ++)
		{
			tot = max(tot, dp_up[x][u] + dp_down[0][k - x][v]);
			if(x != k) tot = max(tot, dp_up[x][u] + dp_down[1][k - x][v] + A[P] - A[u]);
			///printf("x = %d v = %d u = %d cu0 = %lld cu1 = %lld\n", x, v, u, dp_up[x][u] + dp_down[0][k - x][v], dp_up[x][u] + dp_down[1][k - x][v] + A[P] - A[u]);
		}
		for(int x = 0; x <= k; x ++)
		{
			dp_down[0][x][v] = max(dp_down[0][x][v], max(dp_down[0][x][u], dp_down[1][x][u]));

			if(x)
			{
				dp_down[1][x][v] = max(dp_down[1][x][v], nei[v] - A[P] + max(dp_down[1][x - 1][u], dp_down[0][x - 1][u]));
			}
		}
	}
	reverse(all(G[v]));
	for(int i = 0; i < M; i ++) dp_down[0][i][v] = dp_down[1][i][v] = 0;
	for(int i = 1; i <= k; i ++) dp_down[1][i][v] = nei[v] - A[P];
	dp_down[1][0][v] = -inf;
	for(auto u : G[v])
	{
		if(u == P) continue;
		///dfs(u, v);
		for(int x = 0; x <= k; x ++)
		{
			tot = max(tot, dp_up[x][u] + dp_down[0][k - x][v]);
			if(x != k) tot = max(tot, dp_up[x][u] + dp_down[1][k - x][v] + A[P] - A[u]);
		}
		for(int x = 0; x <= k; x ++)
		{
			dp_down[0][x][v] = max(dp_down[0][x][v], max(dp_down[0][x][u], dp_down[1][x][u]));

			if(x)
			{
				dp_down[1][x][v] = max(dp_down[1][x][v], nei[v] - A[P] + max(dp_down[1][x - 1][u], dp_down[0][x - 1][u]));
			}
		}
	}
}

int main()
{
	scanf("%d%d", &n, &k);
    for(int i = 1; i <= n; i ++)
	{
		scanf("%lld", &A[i]);
	}
	for(int i = 1; i < n; i ++)
	{
		int a, b;
		scanf("%d%d", &a, &b);
		G[a].push_back(b);
		G[b].push_back(a);
		nei[a] += A[b];
		nei[b] += A[a];
	}
	dfs_down(1, 0);
	dfs(1, 0);
	for(int i = 1; i <= n; i ++)
	{
		tot = max(tot, dp_down[1][k][i] + A[Par[i]]);
		tot = max(tot, dp_up[k][i]);
	}
	/*for(int i = 1; i <= n; i ++)
	{
	    for(int j = 1; j <= k; j ++)
	    {
		printf("i = %d j = %d dp = %lld\n", i, j, dp_up[j][i]);
	    }
	}
	*/
	printf("%lld", tot);
	return 0;
}

/** test corner cases(n = 1?) watch for overflow or minus indices **/

Compilation message (stderr)

chase.cpp: In function 'int main()':
chase.cpp:105:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  105 |  scanf("%d%d", &n, &k);
      |  ~~~~~^~~~~~~~~~~~~~~~
chase.cpp:108:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  108 |   scanf("%lld", &A[i]);
      |   ~~~~~^~~~~~~~~~~~~~~
chase.cpp:113:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  113 |   scanf("%d%d", &a, &b);
      |   ~~~~~^~~~~~~~~~~~~~~~

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