Submission #778622

#	Time	Username	Problem	Language	Result	Execution time	Memory
778622		1075508020060209tc	Janjetina (COCI21_janjetina)	C++14	110 / 110	524 ms	31352 KiB

Main

#include<bits/stdc++.h>
using namespace std;
#define int long long
#define X first
#define Y second
int lowbit(int x){return x&-x;}
int bit[1000006];
stack<int>stk;
void upd(int pl,int v){
while(pl<=1000000){
    bit[pl]+=v;
    pl+=lowbit(pl);
}
}
int qsum(int pl){
int ret=0;
while(pl){
    ret+=bit[pl];
    pl-=lowbit(pl);
}
return ret;
}
int n;int K;
vector<pair<int,int>>e[500005];
int psmx[500005];
int dph[500005];
int vis[500005];
int sbsz[500005];
int fa[500005];

void dfssz(int nw,int pa){
sbsz[nw]=1;
for(int i=0;i<e[nw].size();i++){
    int v=e[nw][i].first;
    if(v==pa||vis[v]){continue;}
    dfssz(v,nw);
    sbsz[nw]+=sbsz[v];
}
}
void fincen(int nw,int pa,int &cen,int tot){
for(int i=0;i<e[nw].size();i++){
    int v=e[nw][i].first;
    if(v==pa||vis[v]){continue;}
    fincen(v,nw,cen,tot);
}
if(cen==0&&sbsz[nw]*2>=tot){
    cen=nw;
}
}
int ans;
void dfsmx(int nw,int pa){
    fa[nw]=pa;
    dph[nw]=dph[pa]+1;
    for(int i=0;i<e[nw].size();i++){
        int v=e[nw][i].first;int w=e[nw][i].second;
        if(v==pa||vis[v]){continue;}
        w=max(w,psmx[nw]);
        psmx[v]=w;
        dfsmx(v,nw);
    }
}


void solve(int rt){
int cen=0;
dfssz(rt,0);
fincen(rt,0,cen,sbsz[rt]);
psmx[cen]=-1e9;
dfsmx(cen,0);
//cout<<cen<<endl;
priority_queue<pair<int,int>,vector<pair<int,int>>,greater<pair<int,int>>>pq;
for(int i=0;i<e[cen].size();i++){
    int v=e[cen][i].first;
    if(vis[v]){continue;}
    pq.push({psmx[v],v});
}
//upd(1,1);
while(pq.size()){

    int nw=pq.top().second;
  //  cout<<nw<<endl;
    pq.pop();

    int cal=psmx[nw]-(dph[nw]-1)-K;
    if(cal==0){ans++;}
    if(cal>0){
        ans++;
        ans+=qsum(cal+1);
     //   ans-=min(cal,max(0ll,dph[nw]-2ll));
    }
    for(int i=0;i<e[nw].size();i++){
        int v=e[nw][i].first;
        if(v==fa[nw]||vis[v]){continue;}
        pq.push({psmx[v],v});
    }
    upd(dph[nw],1);
    stk.push(dph[nw]);
}
while(stk.size()){
    upd(stk.top(),-1);
    stk.pop();
}

for(int i=0;i<e[cen].size();i++){
    int vv=e[cen][i].first;
    if(vis[vv]){continue;}
    pq.push({psmx[vv],vv});
    while(pq.size()){

        int nw=pq.top().second;
      //  cout<<nw<<endl;
        pq.pop();

        int cal=psmx[nw]-(dph[nw]-1)-K;
        //if(cal==0){ans++;}
        if(cal>0){
            //ans++;
            ans-=qsum(cal+1);
         //   ans-=min(cal,max(0ll,dph[nw]-2ll));
        }
        for(int i=0;i<e[nw].size();i++){
            int v=e[nw][i].first;
            if(v==fa[nw]||vis[v]){continue;}
            pq.push({psmx[v],v});
        }
        upd(dph[nw],1);
        stk.push(dph[nw]);
    }
    while(stk.size()){
        upd(stk.top(),-1);
        stk.pop();
    }
}



vis[cen]=1;
//cout<<cen<<endl;
for(int i=0;i<e[cen].size();i++){
    int v=e[cen][i].first;
    if(vis[v]){continue;}
    solve(v);
}


}



signed main(){
cin>>n>>K;
for(int i=1;i<=n-1;i++){
    int a;int b;int c;
    cin>>a>>b>>c;
    e[a].push_back({b,c});
    e[b].push_back({a,c});
}
ans=0;
solve(1);
cout<<ans*2<<endl;

}

Compilation message (stderr)

Main.cpp: In function 'void dfssz(long long int, long long int)':
Main.cpp:33:14: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   33 | for(int i=0;i<e[nw].size();i++){
      |             ~^~~~~~~~~~~~~
Main.cpp: In function 'void fincen(long long int, long long int, long long int&, long long int)':
Main.cpp:41:14: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   41 | for(int i=0;i<e[nw].size();i++){
      |             ~^~~~~~~~~~~~~
Main.cpp: In function 'void dfsmx(long long int, long long int)':
Main.cpp:54:18: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   54 |     for(int i=0;i<e[nw].size();i++){
      |                 ~^~~~~~~~~~~~~
Main.cpp: In function 'void solve(long long int)':
Main.cpp:72:14: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   72 | for(int i=0;i<e[cen].size();i++){
      |             ~^~~~~~~~~~~~~~
Main.cpp:91:18: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   91 |     for(int i=0;i<e[nw].size();i++){
      |                 ~^~~~~~~~~~~~~
Main.cpp:104:14: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  104 | for(int i=0;i<e[cen].size();i++){
      |             ~^~~~~~~~~~~~~~
Main.cpp:121:22: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  121 |         for(int i=0;i<e[nw].size();i++){
      |                     ~^~~~~~~~~~~~~
Main.cpp:139:14: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  139 | for(int i=0;i<e[cen].size();i++){
      |             ~^~~~~~~~~~~~~~

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