Submission #539114

#	Time	Username	Problem	Language	Result	Execution time	Memory
539114		joelau	Paths (RMI21_paths)	C++14	100 / 100	328 ms	33104 KiB

Main

#include <bits/stdc++.h>
using namespace std;

long long N,K, id[100005], sum = 0, ans[100005], chain[100005];
vector< pair<long long,long long> > lst[100005];
multiset<long long> dp[100005], low, high;

void add (long long x) {
    if (x <= *high.begin()) low.insert(x);
    else {
        high.insert(x);
        sum += x;
        while (high.size() > K) {
            long long v = *high.begin(); high.erase(high.begin());
            sum -= v;
            low.insert(v);
        }
    }
}

void rmv (long long x) {
    if (low.find(x) != low.end()) low.erase(low.find(x));
    else if (high.find(x) != high.end()){
        high.erase(high.find(x));
        sum -= x;
        while (!low.empty() && high.size() < K) {
            long long v = *prev(low.end()); low.erase(prev(low.end()));
            high.insert(v);
            sum += v;
        }
    }
}

void dfs (long long x, long long p) {
    tuple<long long,long long,long long> most = make_tuple(-1,-1,-1);
    for (auto v: lst[x]) if (v.first != p) {
        dfs(v.first,x);
        most = max(most,make_tuple((long long)dp[id[v.first]].size(),v.first,v.second));
    }
    if (most == make_tuple(-1,-1,-1)) {
        id[x] = x;
        dp[x].insert(0);
    }
    else {
        long long a = get<1>(most), b = get<2>(most);
        id[x] = id[a];
        long long tmp = *prev(dp[id[x]].end());
        dp[id[x]].erase(prev(dp[id[x]].end()));
        dp[id[x]].insert(tmp+b);
        chain[a] = tmp+b;
        for (auto v: lst[x]) if (v.first != p && v.first != a) {
            for (auto it = dp[id[v.first]].begin(); it != prev(dp[id[v.first]].end()); ++it) dp[id[x]].insert(*it);
            long long tmp = *prev(dp[id[v.first]].end()) + v.second;
            dp[id[x]].insert(tmp);
            chain[v.first] = tmp;
        }
    }
}

void dfs2 (long long x, long long p, long long pchain) {
    ans[x] = sum;
    long long most = pchain, secondmost = 0;
    for (auto v: lst[x]) if (v.first != p) {
        if (chain[v.first] > most) secondmost = most, most = chain[v.first];
        else if (chain[v.first] > secondmost) secondmost = chain[v.first];
    }
    for (auto v: lst[x]) if (v.first != p) {
        long long nchain = most + v.second;
        if (chain[v.first] == most) nchain = secondmost + v.second;
        add(chain[v.first] - v.second);
        rmv(chain[v.first]);
        add(nchain);
        rmv(nchain - v.second);
        dfs2(v.first,x,nchain);
        add(nchain - v.second);
        rmv(nchain);
        add(chain[v.first]);
        rmv(chain[v.first] - v.second);
    }
}

int main() {
    ios_base::sync_with_stdio(false); cin.tie(NULL);
    cin >> N >> K;
    for (long long i = 0; i < N-1; ++i) {
        long long u,v,w; cin >> u >> v >> w; u--, v--;
        lst[u].emplace_back(v,w), lst[v].emplace_back(u,w);
    }
    dfs(0,-1);
    long long cnt = 0;
    for (auto it = dp[id[0]].begin(); it != dp[id[0]].end(); ++it, ++cnt) {
        if (dp[id[0]].size() - cnt <= K) high.insert(*it), sum += *it;
        else low.insert(*it);
    }
    dfs2(0,-1,0);
    for (long long i = 0; i < N; ++i) cout << ans[i] << '\n';

    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void add(long long int)':
Main.cpp:13:28: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
   13 |         while (high.size() > K) {
      |                ~~~~~~~~~~~~^~~
Main.cpp: In function 'void rmv(long long int)':
Main.cpp:26:44: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
   26 |         while (!low.empty() && high.size() < K) {
      |                                ~~~~~~~~~~~~^~~
Main.cpp: In function 'int main()':
Main.cpp:92:36: warning: comparison of integer expressions of different signedness: 'long long unsigned int' and 'long long int' [-Wsign-compare]
   92 |         if (dp[id[0]].size() - cnt <= K) high.insert(*it), sum += *it;
      |             ~~~~~~~~~~~~~~~~~~~~~~~^~~~

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