Submission #644110

#	Time	Username	Problem	Language	Result	Execution time	Memory
644110		danikoynov	Paths (RMI21_paths)	C++14	100 / 100	467 ms	28440 KiB

Main

/**
 ____ ____ ____ ____ ____ ____
||l |||e |||i |||n |||a |||d ||
||__|||__|||__|||__|||__|||__||
|/__\|/__\|/__\|/__\|/__\|/__\|

**/

#include<bits/stdc++.h>
#define endl '\n'

using namespace std;
typedef long long ll;

void speed()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
}

const ll maxn = 1e5 + 10;
ll n, k, idx[maxn][2], go_up[maxn];
ll num[maxn], depth[maxn], ch[maxn][2], up[maxn], fin[maxn], in[maxn], out[maxn], timer;
vector < pair < ll, ll > > g[maxn];

void calc(ll ver, ll par)
{
    idx[ver][0] = -1;
    idx[ver][1] = -1;
    ch[ver][0] = ch[ver][1] = -1;
    in[ver] = ++ timer;
    for (pair < ll, ll > nb : g[ver])
    {
        if (nb.first == par)
            continue;
        depth[nb.first] = depth[ver] + nb.second;
        calc(nb.first, ver);
        ///cout << ver << " :: " << nb.first << " :: " << ch[nb.first][0] + nb.second << endl;
        if (ch[nb.first][0] + nb.second > ch[ver][0])
        {
            ch[ver][1] = ch[ver][0];
            idx[ver][1] = idx[ver][0];
            ch[ver][0] = ch[nb.first][0] + nb.second;
            idx[ver][0] = idx[nb.first][0];
        }
        else if (ch[nb.first][0] + nb.second > ch[ver][1])
        {
            ch[ver][1] = ch[nb.first][0] + nb.second;
            idx[ver][1] = idx[nb.first][0];
        }
    }

    if (ch[ver][0] == -1)
    {
        ch[ver][0] = 0;
        idx[ver][0] = ver;
    }

    out[ver] = timer;
}

void trav(ll ver, ll par, ll dw)
{
    if (par == -1)
    {
        up[ver] = 0;
        go_up[ver] = 1;
    }
    else
    {
        up[ver] = up[par] + dw;
        go_up[ver] = go_up[par];

        ll path = ch[par][0], go_path = idx[par][0];
        if (idx[par][0] == idx[ver][0])
            path = ch[par][1], go_path = idx[par][1];

        if (path + dw > up[ver])
            up[ver] = path + dw, go_up[ver] = go_path;
    }
    for (pair < ll, ll > nb : g[ver])
    {
        if (nb.first == par)
            continue;
        trav(nb.first, ver, nb.second);
    }
}

void calc_num(ll ver, ll par, ll far)
{
    ll mx = -1;
    bool leaf = true;
    for (pair < ll, ll > nb : g[ver])
    {
        if (nb.first == par)
            continue;
        leaf = false;
        if (idx[nb.first][0] == idx[ver][0])
            calc_num(nb.first, ver, far + nb.second);
        else
            calc_num(nb.first, ver, nb.second);
    }

    if (leaf)
        num[ver] = far;
}

set < pair < ll, ll > > lw, hg;
ll ans = 0;
void add(ll ver, ll val)
{
    pair < ll, ll > cur = {num[ver], ver};
    if (lw.find(cur) != lw.end())
        lw.erase(cur);
    else
    {
        ans -= num[ver];
        hg.erase(cur);
    }

    cur.first += val;
    num[ver] += val;
    hg.insert(cur);
    ///cout << ans << endl;
    ans += cur.first;
    ///cout << cur.first << " :: " << cur.second << endl;
    if (hg.size() > k)
    {
        ans = ans - (*hg.begin()).first;
        lw.insert(*hg.begin());
        hg.erase(hg.begin());
    }

    ///cout << "---------- " << ans << endl;
}

void rem(ll ver, ll val)
{
    pair < ll, ll > cur = {num[ver], ver};
    if (lw.find(cur) != lw.end())
        lw.erase(cur);
    else
    {
        ans -= num[ver];
        hg.erase(cur);
    }

    cur.first -= val;
    num[ver] -= val;
    lw.insert(cur);
    ///cout << "rem " << cur.first << " " << cur.second << endl;
    if (lw.size() > (n - k))
    {
        ans = ans + (*lw.rbegin()).first;
        hg.insert(*lw.rbegin());
        lw.erase(*lw.rbegin());
    }
}

void dfs(ll ver, ll par, ll dw)
{
    //cout << ver << " : " << par << " : " << dw << endl;
    //cout << idx[ver][0] << " " << idx[ver][1] << " " << num[idx[ver][0]] << " " << go_up[ver] << endl;
    if (par != -1)
    {
        if (go_up[ver] > 0)
            add(go_up[ver], dw);
        if (idx[ver][0] > 0)
            rem(idx[ver][0], dw);
    }

    /**cout << "currently" << endl;
    cout << ver << " " << ans << endl;
    for (ll i = 1; i <= n; i ++)
        cout << num[i] << " ";
        cout << endl;*/


    /**ll sum = 0;
    vector < ll > d;
    for (int i = 1; i <= n; i ++)
        d.push_back(num[i]);
    sort(d.begin(), d.end());
    reverse(d.begin(), d.end());
    for (int i = 0; i < k; i ++)
        sum += d[i];
    fin[ver] = sum;*/
    fin[ver] = ans;
    for (pair < ll, ll > nb : g[ver])
    {
        if (nb.first == par)
            continue;
        dfs(nb.first, ver, nb.second);
    }


    if (par != -1)
    {
        if (go_up[ver] > 0)
            rem(go_up[ver], dw);
        if (idx[ver][0] > 0)
            add(idx[ver][0], dw);
    }   /**cout << "finito" << endl;
    cout << ver << endl;
    for (ll i = 1; i <= n; i ++)
        cout << num[i] << " ";
    cout << endl;*/
}
void solve()
{
    cin >> n >> k;
    for (ll i = 1; i < n; i ++)
    {
        ll v, u;
        ll l;
        cin >> v >> u >> l;
        g[v].push_back({u, l});
        g[u].push_back({v, l});
    }

    calc(1, -1);
    trav(1, -1, 0);
    calc_num(1, -1, 0);

    /**for (int i = 1; i <= n; i ++)
    {
        if (go_up[i] <= 0)
            continue;
        if (in[go_up[i]] >= in[i] && out[go_up[i]] <= out[i])
            cout << "FSAWRSF" << endl;
    }*/
    for (ll i = 1; i < k; i ++)
        hg.insert({0, i});
    for (ll i = k; i <= n; i ++)
        lw.insert({0, i});

    for (ll i = 1; i <= n; i ++)
    {
        if (num[i] != 0)
        {
            ll st = num[i];
            num[i] = 0;
            ///cout << "here " << num[i] << " :: " << i << endl;
            add(i, st);
        }
    }

    dfs(1, -1, 0);

    for (ll i = 1; i <= n; i ++)
        cout << fin[i] << endl;
}

int main()
{
    ///freopen("test.txt", "r", stdin);
    ///freopen("output.txt", "w", stdout);
    speed();
    solve();
    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void calc_num(ll, ll, ll)':
Main.cpp:92:8: warning: unused variable 'mx' [-Wunused-variable]
   92 |     ll mx = -1;
      |        ^~
Main.cpp: In function 'void add(ll, ll)':
Main.cpp:128:19: warning: comparison of integer expressions of different signedness: 'std::set<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} and 'll' {aka 'long long int'} [-Wsign-compare]
  128 |     if (hg.size() > k)
      |         ~~~~~~~~~~^~~
Main.cpp: In function 'void rem(ll, ll)':
Main.cpp:153:19: warning: comparison of integer expressions of different signedness: 'std::set<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} and 'll' {aka 'long long int'} [-Wsign-compare]
  153 |     if (lw.size() > (n - k))
      |         ~~~~~~~~~~^~~~~~~~~

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