Submission #639627

#	Time	Username	Problem	Language	Result	Execution time	Memory
639627		Alexandruabcde	Paths (RMI21_paths)	C++14	100 / 100	288 ms	17984 KiB

Main

#include <bits/stdc++.h>

using namespace std;

typedef long long LL;
typedef pair <int, LL> PILL;
typedef pair <LL, int> PLLI;
constexpr int NMAX = 1e5 + 5;

vector <PILL> G[NMAX];
int N, K;
int Root;

int cnt_Leaf;
bool Leaf[NMAX];

LL sol;
LL ans[NMAX];
multiset <LL> Chosen, NotChosen;

void AddPath (LL way) {
    sol += way;
    Chosen.insert(way);

    if (Chosen.size() > K) {
        LL worst = *Chosen.begin();

        sol -= worst;
        Chosen.erase(Chosen.begin());

        NotChosen.insert(worst);
    }
}

void DeletePath (LL way) {
    if (NotChosen.find(way) == NotChosen.end()) {
        if (Chosen.find(way) == Chosen.end())
            return;

        sol -= way;

        Chosen.erase(Chosen.find(way));

        LL best = *prev(NotChosen.end());
        sol += best;

        Chosen.insert(best);
        NotChosen.erase(prev(NotChosen.end()));

        return;
    }

    NotChosen.erase(NotChosen.find(way));
}

void Read () {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> N >> K;

    for (int i = 1; i < N; ++ i ) {
        int x, y, c;

        cin >> x >> y >> c;

        G[x].push_back({y, c});
        G[y].push_back({x, c});
    }

    for (int i = 1; i <= N; ++ i )
        if (G[i].size() == 1) {
            Leaf[i] = true;
            ++ cnt_Leaf;
        }
    for (int i = 1; i <= N; ++ i )
        if (G[i].size() > 1) {
            Root = i;
            break;
        }
}

LL Subtree[NMAX];

void DfsSubtree (int Node, int dad = -1) {
    for (auto it : G[Node]) {
        int to = it.first;
        LL cost = it.second;

        if (to == dad) continue;

        DfsSubtree(to, Node);

        if (Subtree[Node] < Subtree[to] + cost)
            Subtree[Node] = Subtree[to] + cost;
    }
}

LL Ancestor[NMAX];

void DfsAncestor (int Node, int dad = -1) {
    LL Max_1 = 0, Max_2 = 0, who = 0;

    for (auto it : G[Node]) {
        int to = it.first;
        LL cost = it.second;

        if (to == dad) continue;

        if (Max_1 <= Subtree[to] + cost) {
            Max_2 = Max_1;
            Max_1 = Subtree[to] + cost;

            who = to;
        }
        else if (Max_2 < Subtree[to] + cost) Max_2 = Subtree[to] + cost;
    }

    for (auto it : G[Node]) {
        int to = it.first;
        LL cost = it.second;

        if (to == dad) continue;

        Ancestor[to] = max(Ancestor[Node], (to == who ? Max_2 : Max_1)) + cost;

        DfsAncestor(to, Node);

        if (who != to) {
            AddPath(Subtree[to] + cost);
        }
    }
}

void Solve (int Node, int dad = -1) {
    ans[Node] = sol;

    for (auto it : G[Node]) {
        int to = it.first;
        LL cost = it.second;

        if (to == dad) continue;

        DeletePath(Subtree[to] + cost);
        AddPath(Subtree[to]);
        DeletePath(Ancestor[to] - cost);
        AddPath(Ancestor[to]);

        Solve(to, Node);

        AddPath(Subtree[to] + cost);
        DeletePath(Subtree[to]);
        AddPath(Ancestor[to] - cost);
        DeletePath(Ancestor[to]);
    }
}

int main () {
    Read();
    DfsSubtree(Root);
    DfsAncestor(Root);
    AddPath(Subtree[Root]);
    Solve(Root);

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

    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void AddPath(LL)':
Main.cpp:25:23: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   25 |     if (Chosen.size() > 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