oj.uz

Submission #524116

# Submission time Handle Problem Language Result Execution time Memory
524116 2022-02-08T16:29:57 Z wdjpng Paths (RMI21_paths) C++17
100 / 100
 507 ms 37940 KB 
#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

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.0 / 8.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 204 KB Output is correct
2 Correct 1 ms 204 KB Output is correct

Subtask #2
11.0 / 11.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 204 KB Output is correct
2 Correct 1 ms 204 KB Output is correct
3 Correct 1 ms 332 KB Output is correct
4 Correct 1 ms 332 KB Output is correct
5 Correct 1 ms 332 KB Output is correct
6 Correct 1 ms 332 KB Output is correct
7 Correct 1 ms 332 KB Output is correct

Subtask #3
17.0 / 17.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 204 KB Output is correct
2 Correct 1 ms 204 KB Output is correct
3 Correct 1 ms 332 KB Output is correct
4 Correct 1 ms 332 KB Output is correct
5 Correct 1 ms 332 KB Output is correct
6 Correct 1 ms 332 KB Output is correct
7 Correct 1 ms 332 KB Output is correct
8 Correct 3 ms 464 KB Output is correct
9 Correct 3 ms 716 KB Output is correct
10 Correct 3 ms 588 KB Output is correct
11 Correct 4 ms 588 KB Output is correct
12 Correct 3 ms 588 KB Output is correct

Subtask #4
20.0 / 20.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 204 KB Output is correct
2 Correct 1 ms 204 KB Output is correct
3 Correct 1 ms 332 KB Output is correct
4 Correct 1 ms 332 KB Output is correct
5 Correct 1 ms 332 KB Output is correct
6 Correct 1 ms 332 KB Output is correct
7 Correct 1 ms 332 KB Output is correct
8 Correct 3 ms 464 KB Output is correct
9 Correct 3 ms 716 KB Output is correct
10 Correct 3 ms 588 KB Output is correct
11 Correct 4 ms 588 KB Output is correct
12 Correct 3 ms 588 KB Output is correct
13 Correct 6 ms 836 KB Output is correct
14 Correct 13 ms 1072 KB Output is correct
15 Correct 6 ms 972 KB Output is correct
16 Correct 8 ms 864 KB Output is correct
17 Correct 6 ms 844 KB Output is correct
18 Correct 6 ms 716 KB Output is correct
19 Correct 6 ms 844 KB Output is correct

Subtask #5
12.0 / 12.0

# Verdict Execution time Memory Grader output
1 Correct 407 ms 27584 KB Output is correct
2 Correct 397 ms 33568 KB Output is correct
3 Correct 441 ms 27808 KB Output is correct
4 Correct 424 ms 27760 KB Output is correct
5 Correct 389 ms 29816 KB Output is correct
6 Correct 423 ms 27644 KB Output is correct

Subtask #6
32.0 / 32.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 204 KB Output is correct
2 Correct 1 ms 204 KB Output is correct
3 Correct 1 ms 332 KB Output is correct
4 Correct 1 ms 332 KB Output is correct
5 Correct 1 ms 332 KB Output is correct
6 Correct 1 ms 332 KB Output is correct
7 Correct 1 ms 332 KB Output is correct
8 Correct 3 ms 464 KB Output is correct
9 Correct 3 ms 716 KB Output is correct
10 Correct 3 ms 588 KB Output is correct
11 Correct 4 ms 588 KB Output is correct
12 Correct 3 ms 588 KB Output is correct
13 Correct 6 ms 836 KB Output is correct
14 Correct 13 ms 1072 KB Output is correct
15 Correct 6 ms 972 KB Output is correct
16 Correct 8 ms 864 KB Output is correct
17 Correct 6 ms 844 KB Output is correct
18 Correct 6 ms 716 KB Output is correct
19 Correct 6 ms 844 KB Output is correct
20 Correct 407 ms 27584 KB Output is correct
21 Correct 397 ms 33568 KB Output is correct
22 Correct 441 ms 27808 KB Output is correct
23 Correct 424 ms 27760 KB Output is correct
24 Correct 389 ms 29816 KB Output is correct
25 Correct 423 ms 27644 KB Output is correct
26 Correct 479 ms 28228 KB Output is correct
27 Correct 414 ms 34520 KB Output is correct
28 Correct 489 ms 36012 KB Output is correct
29 Correct 404 ms 29120 KB Output is correct
30 Correct 460 ms 29196 KB Output is correct
31 Correct 435 ms 28540 KB Output is correct
32 Correct 483 ms 31696 KB Output is correct
33 Correct 507 ms 29292 KB Output is correct
34 Correct 347 ms 29188 KB Output is correct
35 Correct 399 ms 29244 KB Output is correct
36 Correct 404 ms 37940 KB Output is correct