Submission #1158895

#	Time	Username	Problem	Language	Result	Execution time	Memory
1158895		Mighilon	Petrol stations (CEOI24_stations)	C++20	100 / 100	2837 ms	1043644 KiB

Main

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

 
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;
const int mxN=70005;
vi adj[mxN];
array<int,2> e[mxN];
ll ans[mxN],dist[20][mxN];
int sz[mxN],cnt[mxN],depth[mxN],pred[20][mxN];
bool used[mxN];
map<int,vpi> mp;

struct node{
	ll val;
	node*c[2]={nullptr,nullptr};
	node():val(0){ }
	// ~node(){
	// 	if(c[0]!=nullptr){
	// 		c[0]->~node();
	// 		node *tmp=c[0];
	// 		c[0]=nullptr;
	// 		free(tmp);
	// 	}
	// 	if(c[1]!=nullptr){
	// 		c[1]->~node();
	// 		node *tmp=c[1];
	// 		c[1]=nullptr;
	// 		free(tmp);
	// 	}
	// }
	void push(){
		if(c[0]==nullptr)
			c[0]=new node();
		if(c[1]==nullptr)
			c[1]=new node();
	}
	void merge(){
        if(c[0]&&c[1])
            val=c[0]->val+c[1]->val;
        else if(c[0])
            val=c[0]->val;
        else if(c[1])
            val=c[1]->val;

	}
	void add(ll l, ll r, ll q, ll valq){
		if(q<l||q>r)return;
		if(l==r&&l==q){
			val+=valq;
			return;
		}
		ll m=l+(r-l)/2;
        if(q<=m){
            if(c[0]==nullptr)c[0]=new node();
            c[0]->add(l,m,q,valq);
        }
        else{
            if(c[1]==nullptr)c[1]=new node();
            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;
        ll sum=0;
		ll m=l+(r-l)/2;
        // first optimization
        if(c[0]!=nullptr)
            sum+=c[0]->query(l,m,ql,qr);
        if(c[1]!=nullptr)
            sum+=c[1]->query(m+1,r,ql,qr);
        return sum;
        // 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,ll d,node* tree,int tree_size,int st_sz=0){
	ans[v]+=1LL*(cnt[v]-1)*(tree_size-st_sz);
	trav(id,adj[v]){
		int u=e[id][0]^v;
		if(used[u]||u==p)continue;
		if(!d)trav(x,mp[u]){
			tree->add(0,(n+1)*k,x.f,-x.s);
		}
		ll tmp=tree->query(0,(n+1)*k,d,d+e[id][1]-1);
		ans[v]+=tmp*sz[u];
		tree->add(0,(n+1)*k,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,(n+1)*k,d+k,-tmp);
		if(!d)trav(x,mp[u])tree->add(0,(n+1)*k,x.f,x.s);
	}
}

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

void solve(){
	cin>>n>>k;
	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);
	F0R(i,n)cout<<ans[i]<<nl;
}

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

	int TC = 1;
	// cin >> TC;
    // auto start = chrono::steady_clock::now();
	while(TC--){
		solve();
	}
    // auto end = chrono::steady_clock::now();
    // dbg(chrono::duration_cast<chrono::milliseconds>(end-start).count());
	return 0;
}

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