Submission #784517

#	Time	Username	Problem	Language	Result	Execution time	Memory
784517		blue	Chase (CEOI17_chase)	C++17	100 / 100	505 ms	175128 KiB

chase

#include <bits/stdc++.h>
using namespace std;

using ll = long long;
using vll = vector<ll>;
using vi = vector<int>;
#define sz(x) int(x.size())

const int mx = 100'000;
const int mxv = 100;

int n;
int m;

vll w(1+mx, 0);

vi edge[1+mx];

vi par(1 + mx, 0);

vll wsum(1+mx, 0);

void dfs(int u)
{
	for(int v : edge[u])
	{
		if(par[u] == v)
			continue;

		par[v] = u;
		dfs(v);
	}
}




ll up[1+mx][1+mxv];
ll dn[1+mx][1+mxv];

ll res = 0;

void dfs2(int u)
{
	for(int v : edge[u])
	{
		if(par[u] != v)
			dfs2(v);
	}




	//1. calculate up

	up[u][0] = 0;
	for(int k = 1; k <= m; k++)
		up[u][k] = wsum[u];

	for(int k = 0; k <= m; k++)
	{
		for(int v : edge[u])
		{
			if(par[u] == v)
				continue;

			up[u][k] = max(up[u][k], up[v][k]);
			if(k > 0)
				up[u][k] = max(up[u][k], up[v][k-1] + wsum[u] - w[v]);
		}
	}





	//2. calculate down

	dn[u][0] = 0;
	for(int k = 1; k <= m; k++)
		dn[u][k] = wsum[u] - w[par[u]];

	for(int k = 0; k <= m; k++)
	{
		for(int v : edge[u])
		{
			if(par[u] == v)
				continue;

			dn[u][k] = max(dn[u][k], dn[v][k]);

			if(k > 0)
				dn[u][k] = max(dn[u][k], dn[v][k-1] + wsum[u] - w[par[u]]);
		}
	}



	//3. update res

	//single element
	if(m >= 1)
		res = max(res, wsum[u]);

	//up path
	res = max(res, up[u][m]);

	//down path
	for(int v : edge[u])
	{
		if(par[u] != v)
		{
			res = max(res, dn[v][m]);
			if(m-1 >= 0)
				res = max(res, dn[v][m-1] + wsum[u]);
		}
	}

	//combination

	for(int z = 0; z < 2; z++)
	{
		reverse(edge[u].begin(), edge[u].end());
		vll temp(1+m, -1'000'000'000'000'000'000LL);

		for(int v : edge[u])
		{
			if(par[u] == v)
				continue;

			for(int j = 0; j <= m; j++)
			{
				// res = max(res, temp0[m-j] + dn[v][j]);
				// res = max(res, temp1[m-j] + w[v] + dn[v][j]);
				res = max(res, temp[j] + dn[v][j]);
			}

			////////
			for(int j = 0; j <= m; j++)
			{
				temp[j] = max(temp[j], up[v][m-j]);
				if(m-1-j >= 0)
					temp[j] = max(temp[j], wsum[u] + up[v][m-1-j] - w[v]);
			}
		}
	}

	






	// for(int v1 : edge[u])
	// {
	// 	if(par[u] == v1)
	// 		continue;

	// 	for(int v2 : edge[u])
	// 	{
	// 		if(par[u] == v2)
	// 			continue;

	// 		if(v1 == v2)
	// 			continue;

	// 		for(int j = 0; j <= m; j++)
	// 		{
	// 			res = max(res, (up[v1][m - j]) + (dn[v2][j]));
	// 		}

	// 		for(int j = 0; j <= m-1; j++)
	// 		{
	// 			res = max(res, (wsum[u] + up[v1][m-1 - j] - w[v1]) + (dn[v2][j]));
	// 		}
	// 	}
	// }
}





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

	cin >> n >> m;
	
	for(int i = 1; i <= n; i++)
		cin >> w[i];

	for(int e = 1; e <= n-1; e++)
	{
		int a, b;
		cin >> a >> b;
		edge[a].push_back(b);
		edge[b].push_back(a);
	}

	dfs(1);

	for(int i = 1; i <= n; i++)
		for(int j : edge[i])
			wsum[i] += w[j];

	for(int i = 1; i <= n; i++)
	{
		for(int j = 0; j <= m; j++)
		{
			up[i][j] = 0;
			dn[i][j] = 0;
		}
	}


	dfs2(1);


	// for(int i = 1; i <= n; i++)
	// {
	// 	for(int j = 0; j <= m; j++)
	// 		cerr << dn[i][j] << ' ';
	// 	cerr << '\n';
	// }

	cout << res << '\n';

}

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