Submission #1066253

#	Time	Username	Problem	Language	Result	Execution time	Memory
1066253		MilosMilutinovic	Paths (RMI21_paths)	C++14	100 / 100	221 ms	17804 KiB

Main

#include <bits/stdc++.h>

using namespace std;

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n, k;
  cin >> n >> k;
  vector<vector<pair<int, int>>> g(n);
  for (int i = 1; i < n; i++) {
    int u, v, w;
    cin >> u >> v >> w;
    --u; --v;
    g[u].emplace_back(v, w);
    g[v].emplace_back(u, w);
  }
  multiset<long long> L, R;
  long long s = 0;
  vector<long long> mx(n);
  auto Add = [&](long long x) {
    L.insert(x);
    while ((int) R.size() < k && !L.empty()) {
      long long v = *prev(L.end());
      s += v;
      R.insert(v);
      L.erase(L.find(v));
    }
    while (!L.empty() && !R.empty() && *prev(L.end()) > *R.begin()) {
      long long p = *prev(L.end());
      long long q = *R.begin();
      L.insert(q);
      R.insert(p);
      L.erase(L.find(p));
      R.erase(R.find(q));
      s += p - q;
    }
  };
  auto Remove = [&](long long x) {
    if (R.find(x) != R.end()) {
      s -= x;
      R.erase(R.find(x));
    } else {
      L.erase(L.find(x));
    }
    while (R.size() < k && !L.empty()) {
      long long v = *prev(L.end());
      s += v;
      R.insert(v);
      L.erase(L.find(v));
    }
  };
  function<void(int, int)> Dfs = [&](int v, int pv) {
    for (auto& e : g[v]) {
      int u = e.first;
      int w = e.second;
      if (u == pv) {
        continue;
      }
      Dfs(u, v);
      if (mx[u] != 0) {
        Remove(mx[u]);
      }
      Add(mx[u] + w);
      mx[v] = max(mx[v], mx[u] + w);
    }
  };
  Dfs(0, 0);
  vector<long long> res(n);
  function<void(int, int)> Solve = [&](int v, int pv) {
    long long save = mx[v];
    for (auto& e : g[v]) {
      int u = e.first;
      int w = e.second;
      if (u == pv) {
        if (mx[u] != 0) {
          Remove(mx[u]);
        } 
        Add(mx[u] + w);
        mx[v] = max(mx[v], mx[u] + w);
      }
    }
    res[v] = s;
    pair<long long, int> p = {-1, -1};
    pair<long long, int> q = {-1, -1};
    for (auto& e : g[v]) {
      int u = e.first;
      int w = e.second;
      if (p.first < mx[u] + w) {
        q = p;
        p = {mx[u] + w, u};
      } else if (q.first < mx[u] + w) {
        q = {mx[u] + w, u};
      }
    }
    for (auto& e : g[v]) {
      int u = e.first;
      int w = e.second;
      if (u == pv) {
        continue;
      }
      Remove(mx[u] + w);
      Add(mx[u]);
      if (u != p.second) {
        Solve(u, v);
        Remove(mx[u]);
        Add(mx[u] + w);
        continue;
      }
      long long t = mx[v];
      if (q.second != -1) {
        mx[v] = q.first;
      } else {
        mx[v] = 0;
      }
      Solve(u, v);      
      Remove(mx[u]);
      Add(mx[u] + w);
      mx[v] = t;
    }
    for (auto& e : g[v]) {
      int u = e.first;
      int w = e.second;
      if (u == pv) {
        if (mx[u] != 0) {
          Add(mx[u]);
        } 
        Remove(mx[u] + w);
      }
    }
    mx[v] = save;
  };
  Solve(0, 0);
  for (int i = 0; i < n; i++) {
    cout << res[i] << '\n';
  }
  return 0; 
}

Compilation message (stderr)

Main.cpp: In lambda function:
Main.cpp:46:21: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   46 |     while (R.size() < k && !L.empty()) {
      |            ~~~~~~~~~^~~

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