Submission #377828

#	Time	Username	Problem	Language	Result	Execution time	Memory
377828		kshitij_sodani	Janjetina (COCI21_janjetina)	C++14	110 / 110	584 ms	37480 KiB

Main

//#pragma GCC optimize("Ofast,unroll-loops")
#include <bits/stdc++.h>
using namespace std;
typedef long long llo;
#define mp make_pair
#define pb push_back
#define a first 
#define b second
#define endl '\n'

llo n,k;
vector<pair<llo,llo>> adj[100001];
llo vis[100001];
llo cot[100001];
llo ss[100001];
void dfs(llo no,llo parr=-1){
	ss[no]=1;
	for(auto j:adj[no]){
		if(vis[j.a]==0){
			if(j.a!=parr){
				dfs(j.a,no);
				ss[no]+=ss[j.a];
			}
		}
	}
}
llo pp;
llo find(llo no,llo parr=-1){
	for(auto j:adj[no]){
		if(vis[j.a]==0){
			if(j.a!=parr){
				if(ss[j.a]>ss[pp]/2){
				//	cout<<no<<",,"<<j.a<<endl;
					return find(j.a,no);
				}
			}
		}
	}
	return no;
}
llo lev[100001];
llo par[100001];
vector<llo> tree[100001];
void u(llo ind,llo i,llo j){
	while(i<tree[ind].size()){
		tree[ind][i]+=j;
		i+=(i&-i);
	}
}
llo ss2(llo ind,llo i){
	llo su=0;
	i=min(i,(llo)(tree[ind].size())-1);
	while(i>0){
		su+=tree[ind][i];
		i-=(i&-i);
	}
	return su;
}
vector<llo> zz;
void dfs2(llo no,llo parr=-1,llo levv=0,llo xx=-1){
	if(levv==1){
		xx=no;
		zz.pb(no);
	}
	par[no]=xx;
	lev[no]=levv;
	ss[no]=1;
	for(auto j:adj[no]){
		if(vis[j.a]==0){
			if(j.a!=parr){
				dfs2(j.a,no,levv+1,xx);
				ss[no]+=ss[j.a];
			}
		}
	}
	
}
llo ans=0;
void dec(llo no){
	pp=no;
	dfs(no);
	/*for(int i=0;i<n;i++){
		cout<<ss[i]<<",";
	}
	cout<<endl;*/
	llo cen=find(no);

	dfs2(cen);
	tree[0].clear();
	for(llo i=0;i<=ss[cen];i++){
		tree[0].pb(0);
	}
	for(auto j:zz){
		tree[j+1].clear();
		for(llo m=0;m<=ss[j];m++){
			tree[j+1].pb(0);
		}
	}
	//cout<<cen<<":"<<endl;
	//return ;
	priority_queue<pair<llo,llo>> xx;
	xx.push({0,cen});
	while(xx.size()){
		pair<llo,llo> no=xx.top();
		xx.pop();
		no.a=-no.a;
		//cout<<no.a<<","<<no.b<<endl;
		for(auto j:adj[no.b]){
			if(vis[j.a]==0 and lev[j.a]>lev[no.b]){
				xx.push({-max(no.a,j.b),j.a});
			}
		}
		llo cur=no.a-k-lev[no.b];
		if(cur>=0){
			ans+=ss2(0,cur+1);
			if(no.b!=cen){
				ans-=(ss2(par[no.b]+1,cur+1));
			}
		}

		u(0,lev[no.b]+1,1);
		if(no.b!=cen){
			u(par[no.b]+1,lev[no.b]+1,1);
		}
	}
//	cout<<cen<<":"<<ans<<endl;

	zz.clear();
	vis[cen]=1;
	for(auto j:adj[cen]){
		if(vis[j.a]==0){
			dec(j.a);
		}
	}
}
int main(){
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	cin>>n>>k;
 	for(llo i=0;i<n-1;i++){
 		llo aa,bb,cc;
 		cin>>aa>>bb>>cc;
 		aa--;
 		bb--;
 		adj[aa].pb({bb,cc});
 		adj[bb].pb({aa,cc});
 	}
 	dec(0);
 	ans*=2;
 	cout<<ans<<endl;
	
 
	return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void u(llo, llo, llo)':
Main.cpp:45:9: warning: comparison of integer expressions of different signedness: 'llo' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   45 |  while(i<tree[ind].size()){
      |        ~^~~~~~~~~~~~~~~~~

• Subtask 1: 15 / 15
• Subtask 2: 35 / 35
• Subtask 3: 60 / 60