Submission #1066230

#	Time	Username	Problem	Language	Result	Execution time	Memory
1066230		MilosMilutinovic	Paths (RMI21_paths)	C++14	68 / 100	1083 ms	16712 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> st;
  vector<long long> mx(n);
  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) {
        st.erase(st.find(mx[u]));
      }
      st.insert(mx[u] + w);
      mx[v] = max(mx[v], mx[u] + w);
    }
  };
  Dfs(0, 0);
  auto Calc = [&]() {
    long long s = 0;
    auto it = st.end();
    for (int i = 0; i < k && it != st.begin(); i++) {
      it = prev(it);
      s += *it;
    }
    return s;
  }; 
  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) {
          st.erase(st.find(mx[u]));
        } 
        st.insert(mx[u] + w);
        mx[v] = max(mx[v], mx[u] + w);
      }
    }
    res[v] = Calc();
    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;
      }
      st.erase(st.find(mx[u] + w));
      st.insert(mx[u]);
      if (u != p.second) {
        Solve(u, v);
        st.erase(st.find(mx[u]));
        st.insert(mx[u] + w);
        continue;
      }
      long long t = mx[v];
      if (q.second != -1) {
        mx[v] = q.first;
      } else {
        mx[v] = 0;
      }
      Solve(u, v);      
      st.erase(st.find(mx[u]));
      st.insert(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) {
          st.insert(mx[u]);
        } 
        st.erase(st.find(mx[u] + w));
      }
    }
    mx[v] = save;
  };
  Solve(0, 0);
  for (int i = 0; i < n; i++) {
    cout << res[i] << '\n';
  }
  return 0; 
}

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