Submission #874323

#	Time	Username	Problem	Language	Result	Execution time	Memory
874323		vidut_206_CNH	Chase (CEOI17_chase)	C++14	100 / 100	410 ms	185324 KiB

chase

#include <bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#define pb push_back
#define gcd(a,b) (__gcd(a,b))
#define lcm(a,b) (a/gcd(a,b)*b)
#define sz(x) (int)(x.size())
#define fast_cin() ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
#define db(x) cerr << "[" << "Line " << __LINE__ << " : " << (#x) << " = " << x << "] "

typedef pair<int,int> pii;

const int MOD = 1e9 + 7;
const int MAXN1 = 1e5+5;
const int MAXN2 = 1e6+5;
const long long inf = 1e17;
mt19937_64 rng(time(0));


template<class T>
void maximize(T &res, T val) {
	res = max(res, val);
}

template <class T>
void minimize(T &res, T val) {
	res = min(res, val);
}

const int LOG = 17;

int n,m;
long long w[MAXN1];
long long nei[MAXN1];
int par[MAXN1][LOG + 3];
long long totnei[MAXN1];
long long dist[MAXN1];
int h[MAXN1];
vector<int> adj[MAXN1];

void dfs(int u) {
	dist[u] += w[u];
	totnei[u] += nei[u];
	for(int v : adj[u]) {
		if(v == par[u][0]) continue;
		h[v] = h[u] + 1;
		par[v][0] = u;
		totnei[v] = totnei[u];
		dist[v] = dist[u];
		for(int i = 1; i <= LOG; ++i) par[v][i] = par[par[v][i - 1]][i - 1];
		dfs(v);
	}
}

int lca(int u, int v) {
	if(h[u] < h[v]) swap(u, v);
	int diff = h[u] - h[v];
	for(int i = 0; i <= LOG; ++i) {
		if(diff >> i & 1) u = par[u][i];
	}

	if(u == v) return u;

	for(int i = LOG; i >= 0; --i) {
		if(par[u][i] != par[v][i]) {
			u = par[u][i];
			v = par[v][i];
		}
	}

	return par[u][0];
}

namespace sub1{
	bool check() {
		return (n <= 1000);
	}

	void solve() {
		long long res = 0;

		for(int u = 1; u <= n; ++u) {
			for(int v = 1; v <= n; ++v) {
				int p = lca(u, v);
				long long cost = -w[u];
				long long d = h[u] + h[v] - 2*h[p] + 1;
				if(d > m) continue;

				if(u == v) cost += (nei[u] + w[u]);
				else {
					long long val1 = dist[u] + dist[v] - 2*dist[p] + w[p] - w[u] - w[v];
					long long val2 = totnei[u] + totnei[v] - 2*totnei[p] + nei[p];
					assert(val1 <= val2);
					cost += (val2 - val1);
				}

				maximize(res, cost);
			}
		}

		cout << res;
	}
}

namespace full{

	long long f[MAXN1][103];
	long long g[MAXN1][103];
	pair<long long, int> M[103][3];
	long long res = 0;

	void down(int u) {
		f[u][1] = nei[u];
		for(int v : adj[u]) {
			if(v == par[u][0]) continue;

			down(v);

			for(int i = 1; i <= m; ++i) {
				maximize(f[u][i], f[v][i]);
				if(i > 1) maximize(f[u][i], f[v][i - 1] + nei[u] - w[v]);
				if(i > 1) maximize(f[u][i], f[u][i - 1]);
			}
		}
//		cerr << "\n";
//		db(u);
//		for(int i = 1; i <= m; ++i) db(f[u][i]);
	}
	void ins(int x, long long val, int u) {
		M[x][2] = make_pair(val, u);
		sort(M[x],  M[x] + 3, greater<>());
	}

	void calc(int u) {
		maximize(res, f[u][m]);
		maximize(res, g[u][m]);
//		db(u);
//		cerr << '\n';
//		for(int i = 1; i <= m; ++i) {
//			db(g[u][i]);
//			cerr << "\n";
//		}
//
//		cerr << "\n";

		for(int i = 1; i <= m; ++i) {
			M[i][0] = M[i][1] = M[i][2] = make_pair(-inf, -1);
		}

		for(int i = 1; i < m; ++i) {
			ins(i, g[u][i], u);
		}

		for(int v : adj[u]) {
			if(v == par[u][0]) continue;
			for(int i = 1; i <= m; ++i) {
				long long val = f[v][i] + nei[u] - w[v];
				ins(i + 1, val, v);
				//if(i > 1) ins(i, f[v][i], v);
			}
		}

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

			for(int i = 1; i <= m; ++i) {
				long long val = nei[v] - w[u];

				if(M[i - 1][0].se != v) {
					val += M[i - 1][0].fi;
				}
				else val += M[i - 1][1].fi;

				if(M[i][0].se != v) {
					maximize(g[v][i], M[i][0].fi);
				}
				else maximize(g[v][i], M[i][1].fi);

				if(i > 1) maximize(g[v][i], g[v][i - 1]);
				if(i > 1) maximize(g[v][i], val);
			}
		}

		for(int v : adj[u]) {
			if(v != par[u][0]) calc(v);
		}
	}

	void solve() {
		down(1);

		for(int i = 1; i <= n; ++i) {
			g[i][1] = nei[i];
		}

		calc(1);

		cout << res;
	}
}

int main() {
	fast_cin();


	if(ifstream("tomjerry.inp")) {
		freopen("tomjerry.inp", "r", stdin);
		freopen("tomjerry.out", "w", stdout);
	}
	cin >> n >> m;
	for(int i = 1; i <= n; ++i) cin >> w[i];
	for(int i = 1; i < n; ++i) {
		int u,v;
		cin >> u >> v;
		nei[u] += w[v];
		nei[v] += w[u];
		adj[u].push_back(v);
		adj[v].push_back(u);
	}

	if(m == 0) {
		cout << 0 << "\n";
		exit(0);
	}

	dfs(1);

//	if(sub1::check()) {
//		sub1::solve();
//		return 0;
//	}

	full::solve();




	#ifndef LOCAL
	cerr << "\nTime elapsed: " << 1.0 * (double)clock() / CLOCKS_PER_SEC << " s.\n ";
	#endif

	return 0;
}

Compilation message (stderr)

chase.cpp: In function 'int main()':
chase.cpp:208:10: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  208 |   freopen("tomjerry.inp", "r", stdin);
      |   ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~
chase.cpp:209:10: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  209 |   freopen("tomjerry.out", "w", stdout);
      |   ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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