Submission #1124718

#	Time	Username	Problem	Language	Result	Execution time	Memory
1124718		Zero_OP	Paths (RMI21_paths)	C++17	100 / 100	218 ms	20412 KiB

Main

#include <bits/stdc++.h>

using namespace std;

#define rep(i, l, r) for(int i = (l); i < (r); ++i)
#define sz(v) (int)v.size()
#define dbg(x) "[" #x " = " << (x) << "]"
#define all(v) begin(v), end(v)
#define compact(v) v.erase(unique(all(v)), end(v))
#define file(name) if(fopen(name".inp", "r")){ freopen(name".inp", "r", stdin); freopen(name".out", "w", stdout); }

template<typename T>
    bool minimize(T& a, const T& b){
        if(a > b){
            return a = b, true;
        }
        return false;
    }

template<typename T>
    bool maximize(T& a, const T& b){
        if(a < b){
            return a = b, true;
        }
        return false;
    }

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

using ll = long long;
using ld = long double;
using ull = unsigned long long;
using vi = vector<int>;
using vl = vector<ll>;

const int MAX = 1e5 + 5;

int N, K;
vector<pair<int, int>> adj[MAX];

namespace subtask6{

    pair<ll, int> dp_down[MAX], dp_up[MAX];

    long long sum_k_best = 0, ans[MAX];
    multiset<ll> in_k_best, out_k_best;

    void insert(long long x){
//        cout << "insert : " << x << '\n';
        if(sz(in_k_best) < K){
            in_k_best.insert(x);
            sum_k_best += x;
        } else{
            if(*in_k_best.begin() < x){
                out_k_best.insert(*in_k_best.begin());
                sum_k_best -= *in_k_best.begin();
                in_k_best.erase(in_k_best.begin());
                sum_k_best += x;
                in_k_best.insert(x);
            } else{
                out_k_best.insert(x);
            }
        }

//        cout << '\n';
//        cout << "S_{in} = {";
//        for(auto it : in_k_best) cout << it << ' '; cout << "}\n";
//        cout << "S_{out} = {";
//        for(auto it : out_k_best) cout << it << ' '; cout << "}\n";
    }

    void erase(long long x){
//        cout << "erase :" << x << '\n';
        auto it = in_k_best.lower_bound(x);
        if((it == in_k_best.end()) || (*it != x)){
            auto iter = out_k_best.lower_bound(x);
            assert(iter != out_k_best.end());
            assert(*iter == x);
            out_k_best.erase(iter);
        } else{
            sum_k_best -= x;
            in_k_best.erase(it);

            if(!out_k_best.empty()){
                x = *prev(out_k_best.end());
                out_k_best.erase(prev(out_k_best.end()));
                sum_k_best += x;
                in_k_best.insert(x);
            }
        }
//        cout << '\n';
//        cout << "S_{in} = {";
//        for(auto it : in_k_best) cout << it << ' '; cout << "}\n";
//        cout << "S_{out} = {";
//        for(auto it : out_k_best) cout << it << ' '; cout << "}\n";
    }

    void dfs_down(int u, int p){
        dp_down[u] = {0, u};
        for(auto [v, w] : adj[u]) if(v != p){
            dfs_down(v, u);
            maximize(dp_down[u], make_pair(dp_down[v].first + w, dp_down[v].second));
        }

        for(auto [v, w] : adj[u]) if(v != p){
            if(dp_down[v].second != dp_down[u].second){
                insert(dp_down[v].first + w);
            }
        }

        if(p == -1) insert(dp_down[u].first);
    }

    void dfs_up(int u, int p){
        vector<pair<ll, int>> pref, suff;
        for(auto [v, w] : adj[u]) if(v != p){
            pref.push_back(make_pair(dp_down[v].first + w, dp_down[v].second));
            suff.push_back(make_pair(dp_down[v].first + w, dp_down[v].second));
        }

        for(int i = 1; i < sz(pref); ++i) pref[i] = max(pref[i - 1], pref[i]);
        for(int i = sz(suff) - 2; i >= 0; --i) suff[i] = max(suff[i], suff[i + 1]);

        int i = 0;
        for(auto [v, w] : adj[u]) if(v != p){
            dp_up[v] = {dp_up[u].first + w, dp_up[v].second};
            if(i > 0) maximize(dp_up[v], make_pair(pref[i - 1].first + w, pref[i - 1].second));
            if(i + 1 < sz(suff)) maximize(dp_up[v], make_pair(suff[i + 1].first + w, suff[i + 1].second));
            ++i;
            dfs_up(v, u);
        }
    }

    void dfs_reroot(int u, int p){
        ans[u] = sum_k_best;
//        cout << dbg(u) << dbg(ans[u]) << '\n';
        for(auto [v, w] : adj[u]) if(v != p){
            insert(dp_down[v].first);
            insert(dp_up[v].first);
            erase(dp_down[v].first + w);
            erase(dp_up[v].first - w);
//            cout << '\n';
            dfs_reroot(v, u);
            insert(dp_down[v].first + w);
            insert(dp_up[v].first - w);
            erase(dp_down[v].first);
            erase(dp_up[v].first);
//            cout << '\n';
        }
    }

    void solve(){
        dfs_down(0, -1);
        dfs_up(0, -1);

//        for(int i = 0; i < N; ++i){
//            cout << dbg(i) << dbg(dp_up[i].first) << ' ' << dbg(dp_up[i].second) << '\n';
//        }
//        return;

        if(sz(adj[0]) == 1) insert(0);
        dfs_reroot(0, -1);

        for(int i = 0; i < N; ++i){
            cout << ans[i] << '\n';
        }
    }
}

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    file("B");

    cin >> N >> K;
    for(int i = 1; i < N; ++i){
        int u, v, w;
        cin >> u >> v >> w;
        --u, --v;
        adj[u].emplace_back(v, w);
        adj[v].emplace_back(u, w);
    }

    subtask6::solve();

    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:10:55: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
   10 | #define file(name) if(fopen(name".inp", "r")){ freopen(name".inp", "r", stdin); freopen(name".out", "w", stdout); }
      |                                                ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:174:5: note: in expansion of macro 'file'
  174 |     file("B");
      |     ^~~~
Main.cpp:10:88: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
   10 | #define file(name) if(fopen(name".inp", "r")){ freopen(name".inp", "r", stdin); freopen(name".out", "w", stdout); }
      |                                                                                 ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:174:5: note: in expansion of macro 'file'
  174 |     file("B");
      |     ^~~~

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