Submission #965943

#TimeUsernameProblemLanguageResultExecution timeMemory
965943kilkuwuRoad Closures (APIO21_roads)C++17
22 / 100
188 ms58056 KiB
#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) {

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

  for (int 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<int>> big(N), small(N);
  std::vector<int> 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, int w) {
      if (small[u].empty()) {
        sum_big[u] += w;
        big[u].insert(w);
      } else {
        int 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, int 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<int, 2> {
      vis[u] = dd;
      std::vector<int> to_remove;

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

      std::array<int, 2> res;

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

    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:52:47: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
   52 |       while (small[u].size() && big[u].size() < dd) {
      |                                 ~~~~~~~~~~~~~~^~~~
roads.cpp:59:28: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
   59 |       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, long long int, long long int)> [with auto:23 = minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)::<lambda(auto:23, long long int, long long int)>]':
roads.cpp:146:32:   required from here
roads.cpp:132:25: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
  132 |       if (big[u].size() == dd) {
      |           ~~~~~~~~~~~~~~^~~~~
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...