Submission #1063477

#	Time	Username	Problem	Language	Result	Execution time	Memory
1063477		avighna	Roadside Advertisements (NOI17_roadsideadverts)	C++17	100 / 100	119 ms	27672 KiB

roadsideadverts

#include <bits/stdc++.h>

typedef long long ll;

int main() {
  std::ios_base::sync_with_stdio(false);
  std::cin.tie(nullptr);

  ll n;
  std::cin >> n;
  std::vector<std::vector<std::pair<ll, ll>>> adj(n);
  for (ll i = 0, u, v, w; i < n - 1; ++i) {
    std::cin >> u >> v >> w;
    adj[u].push_back({v, w});
    adj[v].push_back({u, w});
  }

  std::vector<std::vector<ll>> lift(n, std::vector<ll>(16));
  std::vector<std::vector<ll>> lift_sum(n, std::vector<ll>(16));
  std::vector<ll> depth(n);
  std::function<void(ll, ll)> dfs;
  dfs = [&](ll node, ll par) {
    for (auto &[i, w] : adj[node]) {
      if (i == par) {
        continue;
      }
      lift[i][0] = node;
      lift_sum[i][0] = w;
      depth[i] = depth[node] + 1;
      dfs(i, node);
    }
  };
  dfs(0, -1);

  for (ll bt = 1; bt < 16; ++bt) {
    for (ll i = 0; i < n; ++i) {
      lift[i][bt] = lift[lift[i][bt - 1]][bt - 1];
      lift_sum[i][bt] = lift_sum[i][bt - 1] + lift_sum[lift[i][bt - 1]][bt - 1];
    }
  }

  auto lca = [&](ll u, ll v) {
    if (depth[u] > depth[v]) {
      std::swap(u, v);
    }
    ll k = depth[v] - depth[u];
    for (ll bt = 0; bt < 16; ++bt) {
      if (k & (1 << bt)) {
        v = lift[v][bt];
      }
    }
    if (u == v) {
      return u;
    }
    for (ll bt = 15; bt >= 0; --bt) {
      if (lift[u][bt] != lift[v][bt]) {
        u = lift[u][bt];
        v = lift[v][bt];
      }
    }
    return lift[u][0];
  };

  auto lca_mult = [&](const std::vector<ll> &nodes) {
    ll ans = nodes[0];
    for (ll i = 1; i < nodes.size(); ++i) {
      ans = lca(ans, nodes[i]);
    }
    return ans;
  };

  auto sum = [&](ll from, ll to) {
    ll k = depth[from] - depth[to];
    ll ans = 0;
    for (ll bt = 0; bt < 16; ++bt) {
      if (k & (1 << bt)) {
        ans += lift_sum[from][bt];
        from = lift[from][bt];
      }
    }
    return ans;
  };

  ll q;
  std::cin >> q;
  while (q--) {
    std::vector<ll> arr(5);
    std::cin >> arr[0] >> arr[1] >> arr[2] >> arr[3] >> arr[4];
    ll ans = 0;
    while (arr.size() > 1) {
      std::pair<ll, std::pair<ll, ll>> best;
      for (ll i = 0; i < arr.size(); ++i) {
        for (ll j = i + 1; j < arr.size(); ++j) {
          best = std::max(best, {depth[lca(arr[i], arr[j])], {arr[i], arr[j]}});
        }
      }
      auto &[u, v] = best.second;
      auto LCA = lca(u, v);
      arr.erase(std::remove(arr.begin(), arr.end(), u), arr.end());
      arr.erase(std::remove(arr.begin(), arr.end(), v), arr.end());
      arr.push_back(LCA);
      ans += sum(u, LCA);
      ans += sum(v, LCA);
    }
    std::cout << ans << '\n';
  }
}

Compilation message (stderr)

roadsideadverts.cpp: In lambda function:
roadsideadverts.cpp:66:22: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   66 |     for (ll i = 1; i < nodes.size(); ++i) {
      |                    ~~^~~~~~~~~~~~~~
roadsideadverts.cpp: In function 'int main()':
roadsideadverts.cpp:92:24: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   92 |       for (ll i = 0; i < arr.size(); ++i) {
      |                      ~~^~~~~~~~~~~~
roadsideadverts.cpp:93:30: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   93 |         for (ll j = i + 1; j < arr.size(); ++j) {
      |                            ~~^~~~~~~~~~~~
roadsideadverts.cpp:64:8: warning: variable 'lca_mult' set but not used [-Wunused-but-set-variable]
   64 |   auto lca_mult = [&](const std::vector<ll> &nodes) {
      |        ^~~~~~~~

• Subtask 1: 7 / 7
• Subtask 2: 23 / 23
• Subtask 3: 40 / 40
• Subtask 4: 30 / 30