Submission #1087898

#	Time	Username	Problem	Language	Result	Execution time	Memory
1087898		mychecksedad	Petrol stations (CEOI24_stations)	C++17	8 / 100	3553 ms	15444 KiB

Main

#include<bits/stdc++.h>
using namespace std;
#define ll long long int
#define pb push_back
#define en cout << '\n'
#define vi vector<int>
#define pii pair<int, int>
#define ff first
#define ss second
#define all(x) x.begin(),x.end()
const int N = 1e5+100;


int n, k, tin[N], tout[N], timer = 0, par[N], ex[N], compsz[N];
vector<pii> g[N];
ll dp[N], s[N], dist[N], pd[N], PD[N], pdreal[N];
bool is_ancestor(int u, int v){
	return tin[u] <= tin[v] && tout[v] <= tout[u];
}
void f(int v, int p, int up, ll parw){
	// cout << v << ' ' << p << ' ' << up << ' ' << parw << ' ' << dist[v] << "f\n";
	if(dist[v] - dist[up] > k) return;
	if(dist[v] - dist[up] <= k && dist[v] - dist[up] + parw > k){
		// cout << "gh\n";
		dp[up] += dp[v] + compsz[v] - 1;
	}
	for(auto U: g[v]){
		int u = U.ff;
		if(u != p && ex[u] && u != v){
			f(u, v, up, parw);
		}
	}
}
void f2(int v, int p, int up, int up2, ll parw, ll D){
	// cout << v << ' ' << p << ' ' << up << ' ' << parw << ' ' << D << "f\n";
	if(D > k) return;
	// cout << D + parw << ' ' << D << ' ' << parw << '\n';
	if(D <= k && D + parw > k){
		if(is_ancestor(v, up)){
			ll num = 0;
			// cout << v << " : " << up << '\n';
			// ll num = f3(v, p, par[v], dist[p] - dist[v], 0ll);
			pd[up2] += (pd[p] + compsz[v] - 1);
			PD[up2] += (pd[p] + compsz[v] - 1) * s[up2];
		}else{
			// cout << v << ':' << up << '\n';
			pd[up2] += (dp[v] + compsz[v] - 1);
			PD[up2] += (dp[v] + compsz[v] - 1) * s[up2];
		}
	}
	for(auto U: g[v]){
		int u = U.ff;
		// cout << u << ' ' << v << " | ";
		if(u != p && ex[u] && u != v){
			f2(u, v, up, up2, parw, D+U.ss);
		}
	}
	// en;
}

void merg(int a, int b, int pp){
	if(g[a].size() > g[b].size()) g[a].swap(g[b]);
	compsz[b] += compsz[a];
	for(auto x: g[a]) g[b].pb(x);
	if(pp != b) {g[a].swap(g[b]);}
// cout << pp << endl;
// 	for(auto x: g[pp]) cout << x.ff << " "; en;
}
int parww[N];
void dfs(int v, int p, ll parw){
	dp[v] = 1;
	compsz[v] = 1;
	s[v] = 1;
	par[v] = p;
	tin[v] = timer++;
	parww[v] = parw;
	if(v != 1 && parw == 0){
		ex[v] = 0;
		// cout << "merge: " << v << ' ' << p << endl;
		merg(v, p, p);
		return;
	}
	for(int i = 0; i < g[v].size(); ++i){
		auto U = g[v][i];
		int u = U.ff;
		if(u != p && ex[u] && u != v && s[u] == 0){
			// cout << u << ' ' << v << endl;
			dist[u] = dist[v] + U.ss;
			dfs(u, v, U.ss);
			s[v] += s[u];
		}
	}
	tout[v] = timer++;
	f(v, p, v, parw);
}
void dfs2(int v, int p, ll parw){
	pd[v] = 1;
	PD[v] = 0;
	// cout << v << '\n';
	if(p != v){
		f2(p, v, p, v, parw, 0ll);
	}
	for(auto U: g[v]){
		int u = U.ff;
		if(u != p && ex[u] && u != v){
			dfs2(u, v, U.ss);
		}
	}
}
void solve(){
	cin >> n >> k;
	for(int i = 1; i < n; ++i){
		int u, v, w; cin >> u >> v >> w;
		++u, ++v;
		g[u].pb({v, w});
		g[v].pb({u, w});
	}
	for(int i = 1; i <= n; ++i){
		ex[i] = 1;
		s[i] = 0;
	}
	dfs(1, 1, 0);
	// for(int i = 1; i <= n; ++i){
	// 	cout << dp[i] << ' ';
	// }
	for(int i = 2; i <= n; ++i){
		sort(all(g[i]));
		int pos = lower_bound(all(g[i]), pair<int,int>{par[i], -1}) - g[i].begin();
		if(pos == g[i].size() || g[i][pos].ff != par[i]){
			g[i].pb({par[i], parww[i]});
		}
	}
	dfs2(1, 1, 0);
	// for(int i = 1; i <= n; ++i){
	// 	cout << dp[i] << ' ';
	// }
	// for(int i = 1; i <= n; ++i) cout << dp[i] << ' ' << pd[i] << '\n';
	for(int i = 1; i <= n; ++i){
		if(ex[i] == 0){
			cout << "0\n";
			continue;
		}
		pdreal[i] = 0;
		for(auto U: g[i]){
			if(U.ff != par[i] && U.ff != i && ex[U.ff]) pdreal[i] += PD[U.ff];
		}
		cout << (dp[i] - 1) * (n - s[i]) + pdreal[i] << '\n';
	}
}

int main(){
	cin.tie(0); ios::sync_with_stdio(0);
	solve();
	return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void f2(int, int, int, int, long long int, long long int)':
Main.cpp:40:7: warning: unused variable 'num' [-Wunused-variable]
   40 |    ll num = 0;
      |       ^~~
Main.cpp: In function 'void dfs(int, int, long long int)':
Main.cpp:83:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   83 |  for(int i = 0; i < g[v].size(); ++i){
      |                 ~~^~~~~~~~~~~~~
Main.cpp: In function 'void solve()':
Main.cpp:129:10: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  129 |   if(pos == g[i].size() || g[i][pos].ff != par[i]){
      |      ~~~~^~~~~~~~~~~~~~

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