Submission #761121

#	Time	Username	Problem	Language	Result	Execution time	Memory
761121		keta_tsimakuridze	Paths (RMI21_paths)	C++17	100 / 100	194 ms	21184 KiB

Main

#include<bits/stdc++.h>
#define f first
#define s second
#define int long long
#define pii pair<int,int>
using namespace std;
const int N = 2e5 + 5, mod = 1e9 + 7; // !
int t, ans[N], mx[N], sum, k;
vector<pii> V[N];
multiset<int> s1, s2;
void add(int x) {
    s2.insert(x);
    sum += x;
    if(s2.size() > k) {
        sum -= *s2.begin();
        s1.insert(*s2.begin());
        s2.erase(s2.find(*s2.begin()));
    }
}
void rem(int x) {
    if(s2.find(x) != s2.end()) {
        s2.erase(s2.find(x));
        sum -= x;
    } else assert(s1.find(x) != s1.end()), s1.erase(s1.find(x));
    if(s2.size() < k && s1.size()) {
        sum += *--s1.end();
        s2.insert(*--s1.end());
        s1.erase(s1.find(*--s1.end()));
    }
}
void dfs(int u, int p) {
    for(int i = 0; i < V[u].size(); i++) {
        int v = V[u][i].f, c = V[u][i].s;
        if(v == p) continue;
        dfs(v, u);
        mx[u] = max(mx[u], mx[v] + c);
    }
    int F = 0;
    for(int i = 0; i < V[u].size(); i++) {
        int v = V[u][i].f, c = V[u][i].s;
        if(v == p) continue;
        if(mx[u] != mx[v] + c || F) {
            add(mx[v] + c);
        }
        else F = 1;
    }
}
void reroot(int u, int p, int up) {
    ans[u] = sum;
    vector<int> MX = {up, -(int)1e15};
    for(auto x : V[u]) {
        int v = x.f, c = x.s;
        if(v == p) continue;
        if(mx[v] + c > MX[0]) {
            swap(MX[0], MX[1]);
            MX[0] = mx[v] + c;
        } else MX[1] = max(MX[1], mx[v] + c);
    }

    for(auto x : V[u]) {
        int v = x.f, c = x.s;
        if(v == p) continue;
        int X = (mx[v] + c == MX[0] ? MX[1] : MX[0]);
        if(X >= 0) rem(X), add(X + c);
        rem(mx[v] + c); add(mx[v]);
        reroot(v, u, X + c);
        rem(mx[v]), add(mx[v] + c);
        if(X >= 0) rem(X + c), add(X);
    }
}
main(){
    ios_base::sync_with_stdio(false),cin.tie(0),cout.tie(0);
    int n;
    cin >> n >> k;
    for(int i = 2; i <= n; i++) {
        int u, v, c;
        cin >> u >> v >> c;
        V[u].push_back({v, c});
        V[v].push_back({u, c});
    }

    dfs(1, 0);
    add(mx[1]); add(0);
    reroot(1, 0, -0);
    for(int i = 1; i <= n; i++) cout << ans[i] << " ";
 }

Compilation message (stderr)

Main.cpp: In function 'void add(long long int)':
Main.cpp:14:18: 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]
   14 |     if(s2.size() > k) {
      |        ~~~~~~~~~~^~~
Main.cpp: In function 'void rem(long long int)':
Main.cpp:25:18: 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]
   25 |     if(s2.size() < k && s1.size()) {
      |        ~~~~~~~~~~^~~
Main.cpp: In function 'void dfs(long long int, long long int)':
Main.cpp:32:22: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   32 |     for(int i = 0; i < V[u].size(); i++) {
      |                    ~~^~~~~~~~~~~~~
Main.cpp:39:22: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   39 |     for(int i = 0; i < V[u].size(); i++) {
      |                    ~~^~~~~~~~~~~~~
Main.cpp: At global scope:
Main.cpp:71:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
   71 | main(){
      | ^~~~

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