Submission #1158811

#	Time	Username	Problem	Language	Result	Execution time	Memory
1158811		Mighilon	Petrol stations (CEOI24_stations)	C++20	55 / 100	3262 ms	2162688 KiB

Main

#include <bits/stdc++.h>
using namespace std;
 
#ifdef DEBUG
#include "../Library/debug.h"
#else
#define dbg(x...)
#endif

#define int long long
 
typedef long long ll;
typedef long double ld;
typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef vector<int> vi;
typedef vector<bool> vb;
typedef vector<ll> vl;
typedef vector<pi> vpi;
typedef vector<pl> vpl; 
 
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define F0R(i, a) for (int i = 0; i < (a); ++i)
#define FORd(i, a, b) for (int i = (b) - 1; i >= (a); --i)
#define F0Rd(i, a) for (int i = (a) - 1; i >= 0; --i)
#define trav(a, x) for (auto& a : x)
#define f first 
#define s second
#define pb push_back
#define sz(x) (int)(x).size()
#define all(x) x.begin(), x.end()
 
const char nl = '\n';
const int INF = 1e9;
// const ll MOD = 1e9 + 7;
const ll MOD=998244353;

int n;
ll k;
vector<vi> adj;
vector<array<int,2>> e;
vi ans,sz,cnt,depth;
vb used;
vector<vl> dist,pred;
map<int,vpi> mp;

struct node{
	int val;
	node*c[2];
	node():val(0){
		c[0]=c[1]=nullptr;
	}
	void push(){
		if(c[0]==nullptr){
			c[0]=new node();
			c[1]=new node();
		}
	}
	void merge(){
		val=c[0]->val+c[1]->val;
	}
	void add(ll l, ll r, ll q, int valq){
		if(q<l||q>r)return;
		if(l==r&&l==q){
			val+=valq;
			return;
		}
		push();
		ll m=l+(r-l)/2;
		c[0]->add(l,m,q,valq);
		c[1]->add(m+1,r,q,valq);
		merge();
	}
	ll query(ll l,ll r,ll ql,ll qr){
		if(r<ql||l>qr)return 0;
		if(ql<=l&&r<=qr)
			return val;
		push();
		ll m=l+(r-l)/2;
		return c[0]->query(l,m,ql,qr)+c[1]->query(m+1,r,ql,qr);
	}
};


int size_tree(int v,int p=-1){
	sz[v]=1;
	trav(id,adj[v]){
		int u=e[id][0]^v;
		if(used[u]||u==p)continue;
		sz[v]+=size_tree(u,v);
	}
	return sz[v];
}

int find_centroid(int v,int size,int p=-1){
	trav(id,adj[v]){
		int u=e[id][0]^v;
		if(u==p||used[u])continue;
		if(sz[u]>size/2)
			return find_centroid(u,size,v);
	}
	return v;
}
void up_stage(int v,int p,int w,int d,int st){
	depth[v]=d;
	pred[0][v]=p;
	dist[0][v]=w;
	sz[v]=1;
	cnt[v]=1;
	FOR(i,1,20){
		pred[i][v]=pred[i-1][pred[i-1][v]];
		dist[i][v]=dist[i-1][v]+dist[i-1][pred[i-1][v]];
	}
	trav(id,adj[v]){
		int u=e[id][0]^v;
		if(used[u]||u==p)continue;
		up_stage(u,v,e[id][1],d+1,d==0?u:st);
		sz[v]+=sz[u];
	}
	int u=v;
	ll dist_up=0;
	F0Rd(i,20){
		if(dist_up+dist[i][u]<=k){
			dist_up+=dist[i][u];
			u=pred[i][u];
		}
	}
	if(depth[u])cnt[u]+=cnt[v];
	else mp[st].pb({k-dist_up,cnt[v]});
}

void down_stage(int v,int p,int d,node* tree,int tree_size,int st_sz=0){
	// up
	// dbg(tree_size,st_sz)
	dbg(cnt[v]-1,tree_size-st_sz)
	ans[v]+=1LL*(cnt[v]-1)*(tree_size-st_sz);
	dbg(v,ans[v])
	trav(id,adj[v]){
		int u=e[id][0]^v;
		if(used[u]||u==p)continue;
		// dbg(tree->query(0,1e15,0,10))
		if(!d)trav(x,mp[u]){
			tree->add(0,1e15,x.f,-x.s);
			// dbg(x)
		}
		ll tmp=tree->query(0,1e15,d,d+e[id][1]-1);
		// dbg(tree->query(0,1e15,0,10))
		dbg(u,tmp,sz[u],d,d+e[id][1])
		// down
		ans[v]+=tmp*sz[u];
		dbg(v,ans[v])
		tree->add(0,1e15,d+k,tmp);
		if(depth[v]==0)st_sz=sz[u];
		down_stage(u,v,d+e[id][1],tree,tree_size,st_sz);
		tree->add(0,1e15,d+k,-tmp);
		if(!d)trav(x,mp[u])tree->add(0,1e15,x.f,x.s);
	}
}

void centroid_decomposition(int c){
	int v=find_centroid(c,size_tree(c));
	// dbg(v)
	// dbg(ans)
	used[v]=true;
	up_stage(v,v,0,0,-1);
	trav(i,mp){
		// dbg(i.f,i.s)
	}
	node *tree=new node();
	trav(i,mp)
		trav(x,i.s)
			tree->add(0,1e15,x.f,x.s);
	down_stage(v,v,0,tree,sz[v]);
	// dbg(ans)
	mp.clear();
	trav(id,adj[v]){
		int u=e[id][0]^v;
		if(used[u])continue;
		centroid_decomposition(u);
	}
}

void solve(){
	cin>>n>>k;
	adj.resize(n);
	ans.resize(n);
	used.resize(n);
	sz.resize(n);
	e.resize(n-1);
	dist.resize(20,vl(n));
	pred.resize(20,vl(n));
	cnt.resize(n);
	depth.resize(n);
	F0R(i,n-1){
		int a,b;
		cin>>a>>b>>e[i][1];
		e[i][0]=a^b;
		adj[a].pb(i);
		adj[b].pb(i);
	}
	centroid_decomposition(0);
	trav(i,ans)cout<<i<<nl;
}

int32_t main(){
	ios::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);

	int TC = 1;
	// cin >> TC;
	while(TC--){
		solve();
	}
	return 0;
}

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