Submission #524117

#	Submission time	Handle	Problem	Language	Result	Execution time	Memory
524117	2022-02-08T16:32:43 Z	wdjpng	Paths (RMI21_paths)	C++17	Compilation error	0 ms	0 KB

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;
multiset<pair<int, int>> dfs(int v, int p, vector<vector<edge>>& E)
{
    multiset<pair<int, int>> cs;
    cs.insert({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;
        ns.insert(tmp);
        bestofsubtree[e.w]=tmp;
 
        if(cs.size()<ns.size()) swap(cs, ns);
        for(auto x : ns) cs.insert(x);
    }
    
    // while (cs.size()>k+3) cs.erase(cs.begin()); 
    
    return cs;
}
 
 
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];
        uchild.first-=e.c;
        ts.replace(bestofsubtree[e.w], uchild);
        
        // add that edge to biggest path not coming from subtree of child 
 
        pair<int, int>tmp;
        if(bestedge.first==-1||true) 
        {
            auto best_it = (ts.topk.rbegin());
            while (pre[(*best_it).second]>=pre[e.w]&&post[(*best_it).second]<=post[e.w]) {
                ++best_it;
                if(best_it==ts.topk.rend())
                {
                    best_it = ts.other.rbegin();
                    while (pre[(*best_it).second]>=pre[e.w]&&post[(*best_it).second]<=post[e.w])
                        ++best_it;
                    break;
                }
            }
            tmp=*best_it;
        } else tmp=bestedge;
        
        auto mod = tmp;
        mod.first+=e.c;
        ts.replace(tmp, mod);
        
        sol[e.w]=ts.sum;
        if(v!=1&&E[v].size()==2) meta_dfs(e.w,v,ts, E, mod);
        else meta_dfs(e.w,v,ts, E, defau);
        // remove child weight from biggest link
        ts.replace(mod, tmp);
 
        // place biggest link to child in again
        ts.replace(uchild, bestofsubtree[e.w]);
    }
}
 
 
signed main()
{
    int n;
    cin>>n>>k;
    
    bestofsubtree.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, defau);
    for(int i = 1;i <= n; i++) cout<<sol[i]<<"\n";

Compilation message

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)
      |                              ~~~~~~~~~~~^~
Main.cpp: In function 'int main()':
Main.cpp:172:5: warning: this 'for' clause does not guard... [-Wmisleading-indentation]
  172 |     for(int i = 1;i <= n; i++) cout<<sol[i]<<"\n";
      |     ^~~
Main.cpp:172: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'for'
  172 |     for(int i = 1;i <= n; i++) cout<<sol[i]<<"\n";
      | 
Main.cpp:172:50: error: expected '}' at end of input
  172 |     for(int i = 1;i <= n; i++) cout<<sol[i]<<"\n";
      |                                                  ^
Main.cpp:145:1: note: to match this '{'
  145 | {
      | ^