Submission #1364432

#TimeUsernameProblemLanguageResultExecution timeMemory
1364432AvianshPetrol stations (CEOI24_stations)C++20
0 / 100
3597 ms134960 KiB
#include <bits/stdc++.h>

using namespace std;


//sparse+/persistant
//first root is 1 NOT 0
struct segTree{
    struct node{
        int sum,l,r;
    };
    node *st;
    node def;
    int cr = 1;
    segTree(int nodes){
        def.sum=0;
        def.l=0;
        def.r=0;
        st=new node[nodes];
        fill(st,st+nodes,def);
    }
    int cpy(int ind){
        st[++cr]=st[ind];
        return cr;
    }
    void update(int l, int r, int i, int val, int ind){
        if(i<l||i>r){
            assert(0);
        }
        if(l==r){
            st[ind].sum+=val;
            return;
        }
        int mid = (l+r)/2;
        if(i<=mid){
            //go left
            st[ind].l=cpy(st[ind].l);
            update(l,mid,i,val,st[ind].l);
        }
        else{
            st[ind].r=cpy(st[ind].r);
            update(mid+1,r,i,val,st[ind].r);
        }
        st[ind].sum=st[st[ind].l].sum+st[st[ind].r].sum;
    }
    int query(int l, int r, int s, int e, int ind){
        if(e<l||s>r){
            return 0;
        }
        if(st[ind].sum==0){
            return 0;
        }
        if(s<=l&&r<=e){
            return st[ind].sum;
        }
        int mid = (l+r)/2;
        return query(l,mid,s,e,st[ind].l)+query(mid+1,r,s,e,st[ind].r);
    }
} seg(1e7);

const int lg = 20;
const int mxn = 70005;
int ans[mxn];
int up[mxn];
int down[mxn];
int n,k;

void find_sz(int st, vector<array<int,2>> g[], int sz[], bool rem[], int p){
    sz[st]=1;
    for(array<int,2>e:g[st]){
        if(e[0]==p||rem[e[0]]){
            continue;
        }
        find_sz(e[0],g,sz,rem,st);
        sz[st]+=sz[e[0]];
    }
}

int findCentroid(int st, vector<array<int,2>>g[], bool rem[], int sz[]){
    for(array<int,2>e:g[st]){
        if(rem[e[0]])
            continue;
        if(sz[e[0]]>sz[st]/2){
            sz[st]-=sz[e[0]];
            sz[e[0]]+=sz[st];
            return findCentroid(e[0],g,rem,sz);
        }
    }
    return st;
}

void parentFinder(int st, vector<array<int,2>>g[], bool rem[], int par[][lg], int p){
    par[st][0]=p;
    for(int i = 1;i<lg;i++){
        par[st][i]=par[par[st][i-1]][i-1];
    }
    for(array<int,2>e:g[st]){
        if(e[0]==p)
            continue;
        parentFinder(e[0],g,rem,par,st);
    }
}

void toFinder(int st, vector<array<int,2>>g[], bool rem[], int par[][lg], int dep[], int d, int num, int segroot){
    dep[st]=d;
    //num has how many outside current centroid
    up[st]=0;
    for(array<int,2>e:g[st]){
        if(e[0]==par[st][0]||rem[e[0]])
            continue;
        toFinder(e[0],g,rem,par,dep,d+e[1],num,segroot);
    }
    ans[st]+=up[st]*num;
    up[st]++;
    //now pass this to wherever i must refuel
    if(d<=k){
        //no refuels needed so must add to correct place in seg
        seg.update(0,k,k-d,up[st],segroot);
    }
    else{
        //there is an intermediate so pass on up[st] to that
        //first find said intermediate
        int curr = st;
        for(int i = lg-1;i>=0;i--){
            int nx = par[curr][i];
            if(dep[st]-dep[nx]>k){
                //bad
            }
            else{
                //good
                curr=nx;
            }
        }
        //now curr should have refuel point when going up hopefully
        up[curr]+=up[st];
        //done
    }
}

void fromFinder(int st, vector<array<int,2>>g[], bool rem[], int sz[], int par[][lg], int dep[], int segtot, int segcurr, int segprev, int num){
    down[st]=0;
    //must calculate down first here
    //will be calculating down for parent in reality
    //although it will be stored in down[st] instead of down[p]
    //2 cases one is when no refuel after cent and one is when there is refuel after cent
    //Case 1(no refuel) via seg
    //count those deps in range [dep[st],dep[par[st][0]]-1]
    int cnt = seg.query(0,k,dep[par[st][0]],dep[st]-1,segtot)-seg.query(0,k,dep[par[st][0]],dep[st]-1,segcurr)+seg.query(0,k,dep[par[st][0]],dep[st]-1,segprev);
    down[st]+=cnt;
    //case 1 handled, now for case2
    //Case 2(refuel) via par bin jump
    int curr = st;
    for(int i = lg-1;i>=0;i--){
        int nx = par[curr][i];
        if(dep[st]-dep[nx]>k){
            //too far
        }
        else{
            //good
            curr=nx;
        }
    }
    //refuel at parent of this
    down[st]+=(curr!=st) ? down[curr] : 0;

    ans[par[st][0]]+=down[st]*sz[st];
    //added answer
    //now go down
    for(array<int,2>e:g[st]){
        if(rem[e[0]]||par[st][0]==e[0])
            continue;
        fromFinder(e[0],g,rem,sz,par,dep,segtot,segcurr,segprev,num);
    }
}

///do going from centroid seprately

void sep(int st, vector<array<int,2>>g[], bool rem[], int d, int p, int sz[]){
    for(array<int,2>e:g[st]){
        if(e[0]==p||rem[e[0]])
            continue;
        if(d+e[1]>k){
            ans[st]+=sz[e[0]];
            sep(e[0],g,rem,e[1],st,sz);
        }
        else{
            sep(e[0],g,rem,d+e[1],st,sz);
        }
    }
}

void decomp(int st, vector<array<int,2>>g[], bool rem[], int sz[], int par[][lg], int dep[]){
    find_sz(st,g,sz,rem,-1);
    st=findCentroid(st,g,rem,sz);
    //make global segtree and use different roots
    //make parents
    parentFinder(st,g,rem,par,st);
    dep[st]=0;
    down[st]=0;
    up[st]=0;
    ///1: Find ones when going to centroid
    ///   1.1: Also stored fuel left when arriving at st, must be kept seperate for each child therefore persistance
    vector<int>kids;
    vector<int>roots;
    for(array<int,2>e:g[st]){
        if(rem[e[0]])
            continue;
        kids.push_back(e[0]);
        if(roots.size()){
            //prefsum with last
            roots.push_back(seg.cpy(roots.back()));
        }
        else{
            //start new
            roots.push_back(seg.cpy(0));
        }
        toFinder(e[0],g,rem,par,dep,e[1],sz[st]-sz[e[0]],roots.back());
    }
    if(kids.size()==0){
        //endpoint so no new anyway go back
        return;
    }
    ///done part 1(hopefully)
    ///2. Find ones when going from centroid
    for(int i = 0;i<kids.size();i++){
        fromFinder(kids[i],g,rem,sz,par,dep,roots.back(),roots[i],(i ? roots[i-1] : 0),sz[st]-1-sz[kids[i]]);
    }
    //should be done atp
    //do seperate
    sep(st,g,rem,0,-1,sz);
    //decomp further
    rem[st]=1;
    for(int i : kids){
        decomp(i,g,rem,sz,par,dep);
    }
}

signed main(){
    ios::sync_with_stdio(0);
    cin.tie(0);
    cin >> n >> k;
    fill(ans,ans+n,0);
    fill(up,up+n,0);
    fill(down,down+n,0);
    vector<array<int,2>>g[n];
    for(int i = 0;i<n-1;i++){
        int a,b,w;
        cin >> a >> b >> w;
        g[a].push_back({b,w});
        g[b].push_back({a,w});
    }
    bool rem[n];
    fill(rem,rem+n,0);
    int par[n][lg];
    int sz[n];
    int dep[n];
    decomp(0,g,rem,sz,par,dep);
    for(int i = 0;i<n;i++){
        cout << ans[i] << "\n";
    }
    return 0;
}
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...