Submission #1129001

#TimeUsernameProblemLanguageResultExecution timeMemory
1129001yash_9a3bPaths (RMI21_paths)C++20
100 / 100
261 ms19824 KiB
#include "bits/stdc++.h"
#pragma GCC optimize("O3,unroll-loops")
#pragma GCC target("avx,avx2,sse,sse2")
#define fast ios_base::sync_with_stdio(0) , cin.tie(0) , cout.tie(0)
#define endl '\n'
#define int long long
#define f first
#define mp make_pair
#define s second
using namespace std;

//figure out n^2 with bottom up dp
const int N = 1e5 + 5;
int n, k;
vector <pair<int,int> > g[N];
int dp[N], leaf[N], ans[N];
bitset<100001> seen;
multiset <int> top, back;
int top_sum = 0;

void shift(){
    while(top.size() < k && back.size()){
        int ed = *back.rbegin();
        back.erase(back.find(ed));
        top_sum += ed;
        top.insert(ed);
    }
    while(back.size()){
        int ed = *back.rbegin(), st = *top.begin();
        if(ed <= st) break;
        back.erase(back.find(ed));
        top.erase(top.find(st));
        top_sum += ed - st;
        back.insert(st);
        top.insert(ed);
    }
}

void ins(int x){
    back.insert(x);
    shift();
}

void del(int x){
    if(back.find(x) != back.end()){
        back.erase(back.find(x));
    }
    else{
        top_sum -= x;
        top.erase(top.find(x));
    }
    shift();
}

void dfs(int node, int par){
    if(g[node].size() == 1 && par != -1){
        leaf[node] = node;
        dp[node] = 0;
        seen[node] = true;
        return;
    }
    pair<int,int> mx = mp(-1, -1);
    for(pair<int,int> i : g[node]){
        if(i.f == par) continue;
        dfs(i.f, node);
        dp[leaf[i.f]] += i.s;
        mx = max(mx, mp(dp[leaf[i.f]], leaf[i.f]));
    }
    dp[mx.s] = mx.f;
    leaf[node] = mx.s;
}

void get_ans(int node, int par){
    ans[node] = top_sum;
    int bot = leaf[node];
    for(pair<int,int> i : g[node]){
        if(i.f == par) continue;
        if(leaf[i.f] == bot){//currently on the max chain
            pair<int,int> nx = mp(-1, -1);
            for(pair<int,int> j : g[node]){
                if(leaf[j.f] == leaf[i.f]) continue;
                // cout << j.f << ": " << leaf[j.f] << ' ' << dp[leaf[j.f]] << endl;
                nx = max(nx, mp(dp[leaf[j.f]], leaf[j.f]));
            }
            // cout << node << ": " << leaf[node] << ": " << nx.f << " " << nx.s << endl;
            // continue;

            del(dp[bot]);
            del(dp[nx.s]);

            // if(dp[bot] == dp[nx.s] && dp[bot] == 26){
            //     cout << "OVER HERE" << endl;
            // }

            dp[bot] -= i.s;
            dp[nx.s] += i.s;
            leaf[node] = nx.s;
            if(dp[nx.s] > dp[bot]){
                // leaf[node] = nx.s;
                leaf[i.f] = nx.s;
            }

            // if(i.s == 7){
            //     cout << bot << ": " << dp[bot] << " " << dp[nx.s] << ' ' << i.f << " " << leaf[i.f] << endl;
            // }

            ins(dp[bot]);
            ins(dp[nx.s]);

            get_ans(i.f, node);

            del(dp[bot]);
            del(dp[nx.s]);

            dp[bot] += i.s;
            dp[nx.s] -= i.s;

            leaf[node] = leaf[i.f] = bot;

            ins(dp[bot]);
            ins(dp[nx.s]);
        }
        else{//not on the max chain --> "train" as they say in 602 idfk
            del(dp[bot]);
            del(dp[leaf[i.f]]);

            dp[bot] += i.s;
            dp[leaf[i.f]] -= i.s;

            ins(dp[bot]);
            ins(dp[leaf[i.f]]);
            int keep = leaf[i.f];
            if(dp[bot] > dp[leaf[i.f]]) leaf[i.f] = bot;

            get_ans(i.f, node);

            del(dp[bot]);
            del(dp[keep]);

            dp[bot] -= i.s;
            dp[keep] += i.s;
            leaf[i.f] = keep;

            ins(dp[bot]);
            ins(dp[leaf[i.f]]);
        }
    }
}

signed main()
{
    fast;
    cin >> n >> k;
    for(int i = 1; i < n; i++){
        int a, b, x; cin >> a >> b >> x;
        g[a].push_back(mp(b, x));
        g[b].push_back(mp(a, x));
    }
    dfs(1, -1);
    for(int i = 1; i <= n; i++){
        if(leaf[i]) ins(dp[i]);
    }
    get_ans(1, -1);
    for(int i = 1; i <= n; i++) cout << ans[i] << endl;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...