Submission #1077223

#	Time	Username	Problem	Language	Result	Execution time	Memory
1077223		pawned	Petrol stations (CEOI24_stations)	C++17	38 / 100	3563 ms	2097152 KiB

Main

#pragma GCC optimize("O1,O2,O3,Ofast,unroll-loops")

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

#define fi first
#define se second
#define pb push_back
typedef long long ll;
typedef pair<int, int> ii;
typedef vector<int> vi;

const char nl = '\n';

void fastIO() {
	ios::sync_with_stdio(false);
	cin.tie(0);
}

const int MAX = 70005;

ll N, K;

vector<ii> adj[MAX];

vector<ll> subsz(MAX, 0);
vi par(MAX, -1);

void dfs(int n) {
	subsz[n] = 1;
	for (ii x : adj[n]) {
		if (subsz[x.fi] == 0) {
			dfs(x.fi);
			par[x.fi] = n;
			subsz[n] += subsz[x.fi];
		}
	}
}

int main() {
	fastIO();
	cin>>N>>K;
	for (int i = 0; i < N - 1; i++) {
		int u, v, w;	// w <= 1e9, can be int
		cin>>u>>v>>w;
		adj[u].pb({v, w});
		adj[v].pb({u, w});
	}
	dfs(0);
/*	cout<<"subsz: ";
	for (int i = 0; i < N; i++) {
		cout<<subsz[i]<<" ";
	}
	cout<<endl;*/
	vi leaves;
	for (int i = 0; i < N; i++) {
		if (adj[i].size() == 1)
			leaves.pb(i);
	}
	map<ii, int> rem;
	for (int i = 0; i < N; i++) {
		for (ii p : adj[i]) {
			rem[{i, p.fi}] = adj[i].size() - 1;
		}
	}
/*	cout<<"rem: "<<endl;
	for (pair<ii, int> p : rem)
		cout<<"("<<p.fi.fi<<", "<<p.fi.se<<"): "<<p.se<<endl;*/
	vi base(K + 2, 0);
	base[K] = 1;
	base[K + 1] = 0;
	map<ii, vi> info;
	queue<ii> q;
	for (int x : leaves) {
		q.push({x, adj[x][0].fi});
		info[{x, adj[x][0].fi}] = base;
	}
	while (!q.empty()) {
		ii x = q.front();
		// {curr, next}
		q.pop();
		for (ii p : adj[x.se]) {
			// p is {vertex, dist}
			if (p.fi == x.fi)
				continue;
			// subtract 1 from rem
			// if rem becomes 0, push
			ii y = {x.se, p.fi};
			rem[y]--;
			vi toput = base;
			if (rem[y] == 0) {
				q.push(y);
				// calculate info for y
				// find all previous
				for (ii q : adj[x.se]) {
					if (q.fi == p.fi)
						continue;
					vi v = info[{q.fi, x.se}];
					for (int j = q.se; j <= K; j++) {	// need refill
						int arrive = j - q.se;
						if (arrive < p.se) {
							toput[K] += v[j];
							if (q.se == 0) {
								if (j != K)
									toput[K + 1] += v[j];
							}
							else {
								toput[K + 1] += v[j];
							}
						}
						else {
							toput[arrive] += v[j];
						}
					}
				}
			}
			info[y] = toput;
		}
	}
/*	cout<<"info: ";
	for (pair<ii, vi> p : info) {
		cout<<"("<<p.fi.fi<<", "<<p.fi.se<<"): ";
		for (int x : p.se)
			cout<<x<<" ";
		cout<<endl;
	}*/
	vector<ll> ans(N, 0);
	for (pair<ii, vi> p : info) {
		ll pre = p.se[K + 1];	// number of refills, excluding start
		int x = p.fi.fi;
		int y = p.fi.se;
//		cout<<"at "<<x<<" "<<y<<endl;
		// going from x to y
		// ans[x] += (the size of y) times pre
		if (par[y] == x) {	// x is above y
			ans[x] += (subsz[y]) * pre;
//			cout<<"adding above "<<(subsz[y]) * pre<<endl;
		}
		else {	// y is above x
			ans[x] += (N - subsz[x]) * pre;
//			cout<<"adding below "<<(N - subsz[x]) * pre<<endl;
		}
	}
//	cout<<"ANSWER: ";
	for (ll x : ans)
		cout<<x<<" ";
	cout<<endl;
}

• Subtask 0: 0 / 0
• Subtask 1: 18 / 18
• Subtask 2: 8 / 8
• Subtask 3: 0 / 10
• Subtask 4: 12 / 12
• Subtask 5: 0 / 17
• Subtask 6: 0 / 35