Submission #770024

#	Time	Username	Problem	Language	Result	Execution time	Memory
770024		vjudge1	Paths (RMI21_paths)	C++17	48 / 100	1083 ms	18756 KiB

Main

#include <bits/stdc++.h>

using namespace std;

#ifdef LOCAL
#include "C:\GCC\debug.h"
#else
#define debug(...) void(42)
#endif

template <typename T, class F = function<T(const T&, const T&)>>
class SparseTable {
 public:
  int n;
  vector<vector<T>> mat;
  F func;

  SparseTable(const vector<T>& a, const F& f) : func(f) {
    n = static_cast<int>(a.size());
    int max_log = 32 - __builtin_clz(n);
    mat.resize(max_log);
    mat[0] = a;
    for (int j = 1; j < max_log; j++) {
      mat[j].resize(n - (1 << j) + 1);
      for (int i = 0; i <= n - (1 << j); i++) {
        mat[j][i] = func(mat[j - 1][i], mat[j - 1][i + (1 << (j - 1))]);
      }
    }
  }

  T get(int from, int to) const {
    assert(0 <= from && from <= to && to <= n - 1);
    int lg = 32 - __builtin_clz(to - from + 1) - 1;
    return func(mat[lg][from], mat[lg][to - (1 << lg) + 1]);
  }
};

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);

  int n, k;
  cin >> n >> k;

  if (k == 1) {
    vector<vector<array<long long, 2>>> adj(n);

    for (int i = 0; i < n - 1; i++) {
      int u, v, c;
      cin >> u >> v >> c;
      --u, --v;
      adj[u].push_back({v, c});
      adj[v].push_back({u, c});
    }

    vector<vector<long long>> d(n, vector<long long> (2));
    vector<long long> res(n);

    function<void(int, int)> Dfs = [&](int u, int prev) {
      for (auto v : adj[u]) {
        if (v[0] == prev) {
          continue;
        }
        Dfs(v[0], u);
        if (d[v[0]][0] + v[1] > d[u][0]) {
          d[u][1] = d[u][0];
          d[u][0] = d[v[0]][0] + v[1];
        } else if (d[v[0]][0] + v[1] > d[u][1]) {
          d[u][1] = d[v[0]][0] + v[1];
        }
      }
    };

    Dfs(0, -1);

    function<void(int, int, long long)> Dfs2 = [&](int u, int prev, long long c) {
      res[u] = max(c, d[u][0]);

      for (auto v : adj[u]) {
        if (v[0] == prev) {
          continue;
        }
        if (d[v[0]][0] + v[1] == d[u][0]) {
          Dfs2(v[0], u, max(c, d[u][1]) + v[1]);
        } else {
          Dfs2(v[0], u, res[u] + v[1]);
        }
      }
    };

    Dfs2(0, -1, 0);

    for (int root = 0; root < n; root++) {
      cout << res[root] << '\n';  
    }
    return 0;
  }

  vector<vector<array<long long, 2>>> adj(n);

  for (int i = 0; i < n - 1; i++) {
    int u, v, c;
    cin >> u >> v >> c;
    --u, --v;
    adj[u].push_back({v, c});
    adj[v].push_back({u, c});
  }

  vector<long long> to(n);
  vector<int> depth(n);
  vector<int> in(n);
  vector<int> leaves;
  vector<int> euler_tour;

  function<void(int, int, long long)> Dfs = [&](int u, int prev, long long sumAlongPath) {
    to[u] = sumAlongPath;
    bool is_leaf = true;

    in[u] = (int) euler_tour.size();
    euler_tour.push_back(u);

    for (auto v : adj[u]) {
      if (v[0] == prev) {
        continue;
      }
      is_leaf = false;
      depth[v[0]] = depth[u] + 1;
      Dfs(v[0], u, sumAlongPath + v[1]);
      euler_tour.push_back(u);
    }

    if (is_leaf) {
      leaves.push_back(u);
    }
  };

  auto Clear = [&]() {
    to.clear();
    to.resize(n);
    depth.clear();
    depth.resize(n);
    in.clear();
    in.resize(n);
    leaves.clear();
    euler_tour.clear();
  };

  for (int root = 0; root < n; root++) {
    Clear();
    Dfs(root, -1, 0LL);

    SparseTable<int> st(euler_tour, [&](int x, int y) {
      if (depth[x] < depth[y]) {
        return x;
      }
      return y;
    });

    auto lca = [&](int x, int y) {
      if (in[x] > in[y]) {
        swap(x, y);
      }
      return st.get(in[x], in[y]);
    };

    vector<int> current_lca(n, root);

    long long res = 0;

    int len = (int) leaves.size();

    for (int j = 0; j < min(len, k); j++) {
      long long best = 0;
      int toUse = -1;
      for (int i = 0; i < (int) leaves.size(); i++) {
        if (best < to[leaves[i]]) {
          best = to[leaves[i]];
          toUse = leaves[i];
        }
      }
      if (toUse == -1) {
        break;
      }
      res += best;

      vector<int> nleaves;
      for (int i = 0; i < (int) leaves.size(); i++) {
        if (leaves[i] == toUse) {
          continue;
        }
        nleaves.push_back(leaves[i]);
        int nlca = lca(toUse, leaves[i]);
        if (depth[nlca] > depth[current_lca[leaves[i]]]) {
          to[leaves[i]] += to[current_lca[leaves[i]]];
          to[leaves[i]] -= to[nlca];
          current_lca[leaves[i]] = nlca;
        }
      }

      swap(leaves, nleaves);
    }

    cout << res << '\n';
  }
  return 0;
}

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