Submission #1050471

#	Time	Username	Problem	Language	Result	Execution time	Memory
1050471		underwaterkillerwhale	Paths (RMI21_paths)	C++17	0 / 100	54 ms	22228 KiB

Main

#include <bits/stdc++.h>
#define ll long long
#define rep(i,m,n) for(int i=(m); i<=(n); i++)
#define reb(i,m,n) for(int i=(m); i>=(n); i--)
#define pii pair<int,int>
#define pll pair<ll,ll>
#define MP make_pair
#define fs first
#define se second
#define bit(msk, i) ((msk >> i) & 1)
#define iter(id, v) for(auto id : v)
#define SZ(v) (int)v.size()
#define ALL(v) v.begin(),v.end()

using namespace std;

mt19937_64 rd(chrono :: steady_clock :: now ().time_since_epoch().count());
ll Rand (ll l, ll r) { return uniform_int_distribution<ll> (l, r) (rd); }

const int N = 1e5 + 7;
const int Mod = 1e9 + 7;
const int INF = 1e9 + 7;
const ll BASE = 137;


struct Edge {
    ll v, w, idx;
};

int n, K, root;
int nleaf = 0;
vector<Edge> ke[N];
int in[N], out[N];

bool candidate[N];
pair<int,int> par[N];

ll Ans[N];
ll added = 0;

struct Segment_Tree {
    int m;
    pll st[N << 2];
    ll lz[N << 2];

    void init (int n) {
        m = n;
        rep (i, 1, m << 2) st[i].fs = 0, lz[i] = 0;
    }

    void down (int id) {
        rep (i, id << 1, id << 1 | 1) {
            st[i].fs += lz[id];
            lz[i] += lz[id];
        }
        lz[id] = 0;
    }

    void update (int id, int l, int r, int u, int v, ll val) {
        if (l > v || r < u) return;
        if (l >= u && r <= v) {
            lz[id] += val;
            st[id].fs += val;
            return;
        }
        down (id);
        int mid = l + r>> 1;
        update (id << 1, l, mid, u, v, val);
        update (id << 1 | 1, mid + 1, r, u, v, val);
        st[id] = max(st[id << 1], st[id << 1 | 1]);
    }

    void Assign (int id, int l, int r, int pos, int lab, ll val) {
        if (l > pos || r < pos) return;
        if (l == r) {
            st[id].fs = val;
            st[id].se = lab;
            return;
        }
        int mid = l + r>> 1;
        Assign (id << 1, l, mid, pos, lab, val);
        Assign (id << 1 | 1, mid + 1, r, pos, lab, val);
        st[id] = max(st[id << 1], st[id << 1 | 1]);
    }
    void update (int u, int v, ll val) {
        update (1, 1, m, u, v, val);
    }
    void Assign (int pos, int lab, ll val) {
        Assign (1, 1, m, pos, lab, val);
    }
}ST;

int time_dfs = 0;
void pdfs (int u, int p, ll Dist = 0) {
    if (SZ(ke[u]) == 1) nleaf++;
    in[u] = ++time_dfs;
    ST.Assign (in[u], u, Dist);
    iter (&id, ke[u]) {
        ll v = id.v, w = id.w, idx = id.idx;
        if (v != p) {
            pdfs (v, u, Dist + w);
            par[v] = MP(u, w);
        }
    }
    out[u] = time_dfs;
}

void dfs2 (int u, int p) {
    iter (&id, ke[u]) {
        ll v = id.v, w = id.w, idx = id.idx;
        if (v == p) continue;
        if (candidate[v] == 0) Ans[v] = Ans[u] + w;
        else Ans[v] = Ans[u] + added * (SZ(ke[v]) == 1);
        dfs2 (v, u);
    }
}

void solve (int root) {
    rep (i, 1, n) candidate[i] = 0;
    time_dfs = 0;
    nleaf = 0;
    ST.init(n);
    pdfs(root, 0);
    K = min(K, nleaf);
    ll res = 0;
    candidate[root] = 1;
    rep (step, 1, K) {
        int X = ST.st[1].se;
        while (X != root) {
            if (candidate[X] == 1) break;
            candidate[X] = 1;
            res += par[X].se;
            ST.update (in[X], out[X], -par[X].se);
            X = par[X].fs;
        }
    }
    added = ST.st[1].fs;
    Ans[root] = res;
}

void solution () {
    cin >> n >> K;
    rep (i, 1, n - 1) {
        ll u, v, w;
        cin >> u >> v >> w;
        ke[u].push_back({v, w, i});
        ke[v].push_back({u, w, i});
    }
    if (n == 2) {
        rep (i, 1, K) cout << ke[1][0].w <<"\n";
        return;
    }
    if (SZ(ke[1]) == 1) {
        iter (&id, ke[1]) {
            if (SZ(ke[id.v]) > 1) {
                root = id.v;
            }
        }
    }
    else root = 1;
    solve (root);
    while (1) {
        int nC = 0, to = -1;
        iter (&id, ke[root]) {
            ll v = id.v, w = id.w, idx = id.idx;
            if (v == par[root].fs) continue;
            if (candidate[v] == 1) ++nC, to = v;
        }
        if (nC == 1) root = to;
        else break;
    }
    solve(root);
    dfs2(root, 0);
    rep (i, 1, n) cout << Ans[i] <<"\n";
}

#define file(name) freopen(name".inp","r",stdin); \
freopen(name".out","w",stdout);
int main () {
//    file("c");
    ios_base :: sync_with_stdio(false); cin.tie(0); cout.tie(0);
    int num_Test = 1;
//    cin >> num_Test;
    while (num_Test--)
        solution();
}
/*
no bug challenge +2
2 + 8 * 2 - 9

*/

Compilation message (stderr)

Main.cpp: In member function 'void Segment_Tree::update(int, int, int, int, int, long long int)':
Main.cpp:67:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   67 |         int mid = l + r>> 1;
      |                   ~~^~~
Main.cpp: In member function 'void Segment_Tree::Assign(int, int, int, int, int, long long int)':
Main.cpp:80:21: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   80 |         int mid = l + r>> 1;
      |                   ~~^~~
Main.cpp: In function 'void pdfs(int, int, long long int)':
Main.cpp:99:32: warning: unused variable 'idx' [-Wunused-variable]
   99 |         ll v = id.v, w = id.w, idx = id.idx;
      |                                ^~~
Main.cpp: In function 'void dfs2(int, int)':
Main.cpp:110:32: warning: unused variable 'idx' [-Wunused-variable]
  110 |         ll v = id.v, w = id.w, idx = id.idx;
      |                                ^~~
Main.cpp: In function 'void solution()':
Main.cpp:165:26: warning: unused variable 'w' [-Wunused-variable]
  165 |             ll v = id.v, w = id.w, idx = id.idx;
      |                          ^
Main.cpp:165:36: warning: unused variable 'idx' [-Wunused-variable]
  165 |             ll v = id.v, w = id.w, idx = id.idx;
      |                                    ^~~

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