Submission #965994

#	Time	Username	Problem	Language	Result	Execution time	Memory
965994		kilkuwu	Road Closures (APIO21_roads)	C++17	100 / 100	208 ms	54480 KiB

roads

#include "roads.h"

#include <bits/stdc++.h>

#ifdef LOCAL
#include "template/debug.hpp"
#else
#define dbg(...) ;
#define timer(...) ;
#endif

std::vector<long long> minimum_closure_costs(int N, std::vector<int> U,
                                             std::vector<int> V,
                                             std::vector<int> W) {

  std::vector<long long> res(N);
  std::vector<std::vector<std::pair<int, int>>> adj(N);
  long long sum_edges_weight = 0;

  for (long long i = 0; i + 1 < N; i++) {
    int u = U[i];
    int v = V[i];
    int w = W[i];
    adj[u].emplace_back(v, w);
    adj[v].emplace_back(u, w); 
    sum_edges_weight += w;
  }

  std::vector<std::vector<int>> at_deg(N);
  for (int i = 0; i < N; i++) {
    at_deg[adj[i].size()].push_back(i);
  }

  // now what do we do 
  res[0] = sum_edges_weight;

  auto not_considered = sum_edges_weight;
  std::vector<int> added(N);

  std::vector<int> vertices;

  std::vector<std::vector<std::pair<int, int>>> g(N);

  std::vector<int> vis(N, -1);

  std::vector<std::multiset<long long>> big(N), small(N);
  std::vector<long long> sum_big(N);

  for (int dd = N - 1; dd > 0; dd--) {
    auto balance = [&](int u) {
      while (small[u].size() && big[u].size() < dd) {
        auto it = --small[u].end();
        big[u].insert(*it);
        sum_big[u] += *it;
        small[u].erase(it);
      }
      
      while (big[u].size() > dd) {
        auto it = big[u].begin();
        small[u].insert(*it);
        sum_big[u] -= *it;
        big[u].erase(it);
      }
    };

    auto add_leaf = [&](int u, long long w) {
      if (small[u].empty()) {
        sum_big[u] += w;
        big[u].insert(w);
      } else {
        long long m = *(--small[u].end());
        if (w <= m) {
          small[u].insert(w);
        } else {
          sum_big[u] += w;
          big[u].insert(w);
        }
      }
      balance(u);
    };

    auto remove_leaf = [&](int u, long long w) {
      auto it = small[u].find(w);
      if (it == small[u].end()) {
        sum_big[u] -= w;
        big[u].erase(big[u].find(w));
      } else {
        small[u].erase(it);
      }
      balance(u);
    };
    
    for (int u : at_deg[dd]) {
      added[u] = 1;
      vertices.push_back(u);
      for (auto [v, w] : adj[u]) {
        if (added[v]) {
          g[u].emplace_back(v, w);
          g[v].emplace_back(u, w); 
          remove_leaf(v, w);
        } else {
          not_considered -= w;
          add_leaf(u, w);
        }
      }
    }

    for (int v : vertices) {
      balance(v);
      // dbg(dd, v, big[v]);
    }

    auto dfs = [&](auto self, int u, int p) -> std::array<long long, 2> {
      vis[u] = dd;
      std::vector<long long> to_remove;

      auto remove_all = 0LL;
      for (auto [v, w] : g[u]) {
        if (v == p) continue;
        auto rr = self(self, v, u);
        remove_all += rr[0];
        auto delta = rr[1] + w - rr[0];
        if (delta > 0) {
          to_remove.push_back(delta);
          add_leaf(u, delta);
        }
      }

      std::array<long long, 2> rr;

      rr[0] = sum_big[u] + remove_all;
      rr[1] = rr[0];
      if (big[u].size() == dd) {
        rr[1] -= *big[u].begin();
      }
      
      for (int x : to_remove) {
        remove_leaf(u, x);
      }
      // dbg(dd, u, p, rr);
      return rr;
    };

    auto saved = 0LL;

    for (int u : vertices) {
      if (vis[u] == dd) continue;
      auto res = dfs(dfs, u, -1);
      saved += res[0]; // nothing is connected to it 
    }

    res[dd] = sum_edges_weight - (not_considered + saved);
  }

  return res;
}

Compilation message (stderr)

roads.cpp: In lambda function:
roads.cpp:51:47: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   51 |       while (small[u].size() && big[u].size() < dd) {
      |                                 ~~~~~~~~~~~~~~^~~~
roads.cpp:58:28: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   58 |       while (big[u].size() > dd) {
      |              ~~~~~~~~~~~~~~^~~~
roads.cpp: In instantiation of 'minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)::<lambda(auto:23, int, int)> [with auto:23 = minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)::<lambda(auto:23, int, int)>]':
roads.cpp:148:32:   required from here
roads.cpp:133:25: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  133 |       if (big[u].size() == dd) {
      |           ~~~~~~~~~~~~~~^~~~~

• Subtask 1: 5 / 5
• Subtask 2: 7 / 7
• Subtask 3: 14 / 14
• Subtask 4: 10 / 10
• Subtask 5: 17 / 17
• Subtask 6: 25 / 25
• Subtask 7: 22 / 22