Submission #644106

#	Time	Username	Problem	Language	Result	Execution time	Memory
644106		danikoynov	Paths (RMI21_paths)	C++14	31 / 100	346 ms	24776 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 int maxn = 1e5 + 10;
int n, k, idx[maxn][2], go_up[maxn];
ll num[maxn], depth[maxn], ch[maxn][2], up[maxn], fin[maxn];
vector < pair < int, ll > > g[maxn];

void calc(int ver, int par)
{
    idx[ver][0] = -1;
    idx[ver][1] = -1;
    ch[ver][0] = ch[ver][1] = -1;

    for (pair < int, 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;
    }

}

void trav(int ver, int 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 < int, ll > nb : g[ver])
    {
        if (nb.first == par)
            continue;
        trav(nb.first, ver, nb.second);
    }
}

void calc_num(int ver, int par, ll far)
{
    int mx = -1;
    bool leaf = true;
    for (pair < int, 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, int > > lw, hg;
ll ans = 0;
void add(int ver, ll val)
{
    pair < ll, int > 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);
    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());
    }

}

void rem(int ver, ll val)
{
    pair < ll, int > 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);

    if (lw.size() > (n - k))
    {
        ans = ans + (*lw.rbegin()).first;
        hg.insert(*lw.rbegin());
        lw.erase(*lw.rbegin());
    }
}

void dfs(int ver, int par, ll dw)
{
    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 << endl;
    for (int i = 1; i <= n; i ++)
        cout << num[i] << " ";
        cout << endl;*/
    fin[ver] = ans;
    for (pair < int, 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 (int i = 1; i <= n; i ++)
        cout << num[i] << " ";
    cout << endl;*/
}
void solve()
{
    cin >> n >> k;
    for (int i = 1; i < n; i ++)
    {
        int 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 <= k; i ++)
    hg.insert({0, i});
    for (int i = k + 1; i <= n; i ++)
    lw.insert({0, i});

    for (int 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 (int i = 1; i <= n; i ++)
        cout << fin[i] << endl;
}

int main()
{
    speed();
    solve();
    return 0;
}

Compilation message (stderr)

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

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