제출 #1242353

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1242353		ender_shayan	Janjetina (COCI21_janjetina)	C++20	110 / 110	271 ms	70628 KiB

Main

#include <bits/stdc++.h>

using namespace std;

// #include <ext/pb_ds/assoc_container.hpp>
// #include <ext/pb_ds/tree_policy.hpp>
// using namespace __gnu_pbds;

typedef long long ll;
typedef long double	ld;
typedef pair<int, int>	pii  ;
typedef pair<ll, ll>	pll  ;
typedef vector<pii>     vii  ;
typedef vector<int>     veci ;
typedef vector<pll>     vll  ;
typedef vector<ll>      vecll;

// find_by_order             order_of_key

//#pragma GCC optimize("O3,unroll-loops")
//#pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#define ordered_set tree<int, null_type,less<int>, rb_tree_tag,tree_order_statistics_node_update>
#define F		        first
#define S		        second
#define pb		        push_back
#define endl            '\n'
#define Mp		        make_pair
#define all(x)          x.begin(), x.end()
#define debug(x)        cerr << #x << " = " << x << endl
#define set_dec(x)	    cout << fixed << setprecision(x);
#define fast_io         ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
#define file_io         freopen("in.txt" , "r" , stdin) ; freopen("out.txt" , "w" , stdout);
#define lb              lower_bound
#define ub              upper_bound
#define for1(n)         for(int i=1;i<=n;i++)
#define for0(n)         for(int i=0;i<n;i++)
#define forn(n)         for(int i=n;i>0;i--)
#define pq              priority_queue <pii, vector<pii>, greater<pii>>


const int N=1e6+10;
int n,m,k,q,vis[N],sz[N],fen[N],nn,mol[N];
vector<pii>g[N];
vector<pair<pii,int>>vec;
vector<int> vecc[N];
ll ans;
void upd(int i,int x){
    for(i++;i<N;i+=i&-i)fen[i]+=x;
}
int get(int i){
    int ans=0;
    for(i++;i>0;i-=i&-i)ans+=fen[i];
    return ans;
}
void upd(int t,int i,int x){
    for(i++;i<vecc[t].size();i+=i&-i)vecc[t][i]+=x;
}
int get(int t,int i){
    int ans=0;
    for(i++;i>0;i-=i&-i)ans+=vecc[t][i];
    return ans;
}
void DFS(int v,int p=0){
    sz[v]=1;
    for(auto [u,w]:g[v])if(u!=p && vis[u]==0){
        DFS(u,v);
        sz[v]+=sz[u];
    }
}
int cen(int v,int p=0){
    for(auto[u,w]:g[v])if(vis[u]==0 && u!=p && sz[u]>=nn/2)return cen(u,v);
    return v;
}
void dfs(int v,int p=0,int l=0,int w=0){
    vec.pb({{w,l},mol[v]});
    while(vecc[mol[v]].size()<=l+1)vecc[mol[v]].pb(0);
    for(auto [u,ww]:g[v])if(u!=p && vis[u]==0){
        if(p==0)mol[u]=u;
        else mol[u]=mol[v];
        dfs(u,v,l+1,max(w,ww));
    }
}
void solve(int root){
    DFS(root);nn=sz[root];
    root=cen(root);mol[root]=root;
    dfs(root);
    sort(all(vec));
    for(auto[p,m]:vec){
        int w=p.F,l=p.S;
        ans+=get(w-l-k)-get(m,min(w-l-k,int(vecc[m].size())-2));
        upd(l,1);
        upd(m,l,1);
    }
    for(auto[p,m]:vec)upd(p.S,-1);

    vec.clear();vis[root]=2;vecc[root].clear();
    for(auto [u,w]:g[root])if(vis[u]==0){vecc[u].clear();solve(u);}
}
int main(){
    fast_io
    cin>>n>>k;
    for1(n-1){
        int u,v,w;cin>>u>>v>>w;
        g[u].pb({v,w});
        g[v].pb({u,w});
    }
    solve(1);
    cout<<ans*2<<endl;




}

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