Submission #524116

#	Time	Username	Problem	Language	Result	Execution time	Memory
524116		wdjpng	Paths (RMI21_paths)	C++17	100 / 100	507 ms	37940 KiB

Main

#ifndef LOCAL
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
#endif

#include <bits/stdc++.h>

#define int long long
#define rep(i,n) for(int i =0; i<n;i++)
#define all(a) a.begin(), a.end()

using namespace std;

int k;
struct edge{
    int c, w;
};

struct topkmultiset{
    int k;
    int sum = 0;
    multiset<pair<int, int>>other,topk;
    void rebalance()
    {
        while (other.size()&&topk.size()<k)
        {
            auto a = other.rbegin();
            auto end  = *a;
            sum+=end.first;
            topk.insert(end);
            other.erase(end);
        }
        
        while (other.size()&&topk.size()&&*(other.rbegin())>*(topk.begin()))
        {
           auto a = *(other.rbegin()), b = *(topk.begin());
           sum+=a.first-b.first;
           topk.erase(b);
           other.insert(b);
           
           other.erase(a);
           topk.insert(a);
        }
        
    }

    void replace(pair<int, int> a, pair<int, int> b) {
        if(topk.find(a)!=topk.end()) {
            sum-=a.first;
            topk.erase(a);
        } else other.erase(a);

        other.insert(b);
        rebalance();
    }
};

int precount=0;
vector<int>pre,post;
void predfs(int v, int p, vector<vector<edge>>&E)
{
    pre[v]=precount++;
    for(auto e : E[v]) if(e.w!=p) predfs(e.w,v,E);
    post[v]=precount++;
}

vector<pair<int, int>>bestofsubtree;
vector<pair<int, int>>bestofsubtreefromotherchild;
multiset<pair<int, int>> dfs(int v, int p, vector<vector<edge>>& E)
{
    multiset<pair<int, int>> cs;
    cs.insert({0,v});

    pair<int, int>best={0,v}, second_best = {0,v};

    for(edge e : E[v])
    {
        if(e.w==p) continue;
        multiset<pair<int, int>> ns = dfs(e.w, v, E);
        
        auto it = ns.end();
        it--;
        pair<int,int> tmp = *it;

        ns.erase(it);
        tmp.first+=e.c;

        if(tmp>best) {second_best=best;best=tmp;}
        else second_best=max(second_best, tmp);

        ns.insert(tmp);

        if(cs.size()<ns.size()) swap(cs, ns);
        for(auto x : ns) cs .insert(x);
    }
    bestofsubtreefromotherchild[v] = second_best;
    bestofsubtree[v] = best;
    // while (cs.size()>k+3) cs.erase(cs.begin()); 
    
    return cs;
}


bool ischild(int w, int v) {return pre[v]<=pre[w]&&post[w]<=post[v];}
vector<int>sol;
pair<int, int>defau = {-1,-1};
void meta_dfs(int v, int p, topkmultiset&ts, vector<vector<edge>>& E, pair<int, int>bestedge)
{
    for(edge e : E[v])
    {
        if(e.w==p) continue;

        // remove edge from child to root from biggest path of child
        pair<int, int>uchild = bestofsubtree[e.w];
        pair<int, int>mchid = uchild;
        mchid.first+=e.c;
        ts.replace(mchid, uchild);
        
        // add that edge to biggest path not coming from subtree of child 

        pair<int, int>tmp;
        tmp = bestofsubtree[v];
        if(ischild(tmp.second, e.w) || (bestofsubtreefromotherchild[v] > tmp &&(!ischild(bestofsubtreefromotherchild[v].second, e.w)))) 
            tmp = bestofsubtreefromotherchild[v];
        if(bestedge > tmp&&!(ischild(bestedge.second, e.w))) tmp = bestedge;

        auto mod = tmp;
        mod.first+=e.c;
        ts.replace(tmp, mod);

        sol[e.w]=ts.sum;
        meta_dfs(e.w,v,ts, E, mod);
       
        // remove child weight from biggest link
        ts.replace(mod, tmp);

        // place biggest link to child in again
        ts.replace(uchild, mchid);
    }
}


signed main()
{
    int n;
    cin>>n>>k;
    
    bestofsubtree.assign(n+1, {-1,-1});
    bestofsubtreefromotherchild.assign(n+1, {-1,-1});
    pre.assign(n+1, -1);
    post.assign(n+1,-1);

    vector<vector<edge>>E(n+1);
    rep(i, n-1)
    {
        int a, b, c;
        cin>>a>>b>>c;

        E[a].push_back({c,b});
        E[b].push_back({c,a});
    }
    predfs(1,-1,E);
    sol.resize(n+1);
    multiset<pair<int, int>>s  = dfs(1, -1, E);
    topkmultiset ts = topkmultiset();
    ts.k=k;
    ts.other = s;
    ts.rebalance();
    sol[1]=ts.sum;

    meta_dfs(1, -1, ts, E, {0,0});
    for(int i = 1;i <= n; i++) cout<<sol[i]<<"\n";
}

Compilation message (stderr)

Main.cpp: In member function 'void topkmultiset::rebalance()':
Main.cpp:25:41: warning: comparison of integer expressions of different signedness: 'std::multiset<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
   25 |         while (other.size()&&topk.size()<k)
      |                              ~~~~~~~~~~~^~

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