Submission #1128997

#	Time	Username	Problem	Language	Result	Execution time	Memory
1128997		yash_9a3b	Paths (RMI21_paths)	C++20	0 / 100	203 ms	16012 KiB

Main

#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;
// vector<int> changed;
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){
    // changed.push_back(x);
    if(back.find(x) != back.end()){
        back.erase(back.find(x));
    }
    else{
        top_sum -= x;
        // if(top.count(x) == 0){
        //     cout << "here: ";
        //     for(int i : changed) cout << i << " ";
        //     cout << endl;
        //     changed.clear();
        //     return;
        // }
        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 << 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;
            if(dp[nx.s] > dp[bot]){
                leaf[node] = nx.s;
                leaf[i.f] = nx.s;
            }

            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;
}

• Subtask 1: 0 / 8
• Subtask 2: 0 / 11
• Subtask 3: 0 / 17
• Subtask 4: 0 / 20
• Subtask 5: 0 / 12
• Subtask 6: 0 / 32