Submission #219759

#	Time	Username	Problem	Language	Result	Execution time	Memory
219759		rama_pang	Designated Cities (JOI19_designated_cities)	C++14	100 / 100	1316 ms	36056 KiB

designated_cities

#include <bits/stdc++.h>
using namespace std;
using lint = long long;

struct edge_t {
  int v;
  int away;
  int come;
  edge_t() {}
  edge_t(int v, int a, int c) : v(v), away(a), come(c) {}
};

int N;
vector<vector<edge_t>> adj;

vector<lint> Answer;
vector<lint> BaseCost; // BaseCost[i] = sum of all edges going to i (root at i), so we do not need to consider them afterwards
vector<bool> processed; // processed as R, done via centroid decomposition

vector<int> sz; // subtree size
vector<lint> optimal; // maximum weighted depth

void GenerateBaseCost() {
  vector<lint> dp(N);
  function<void(int, int)> dfs = [&](int n, int p) {
    for (const auto &i : adj[n]) if (i.v != p) {
      dfs(i.v, n);
      dp[n] += dp[i.v] + i.come;
    }
  };

  function<void(int, int)> reroot = [&](int n, int p) {
    BaseCost[n] = dp[n];
    for (const auto &i : adj[n]) if (i.v != p) {
      dp[n] -= dp[i.v] + i.come;
      dp[i.v] += dp[n] + i.away;
      reroot(i.v, n);
      dp[i.v] -= dp[n] + i.away;
      dp[n] += dp[i.v] + i.come;
    }
  };

  dfs(0, -1);
  reroot(0, -1);

}

int Centroid(int n) {
  function<void(int, int)> dfs = [&](int n, int p) {
    sz[n] = 1;
    for (const auto &i : adj[n]) if (i.v != p && !processed[i.v]) {
      dfs(i.v, n);
      sz[n] += sz[i.v];
    }
  };

  dfs(n, -1);
  int cur = n, p = -1;
  while (true) {
    pair<int, int> mx = {-1, -1};
    for (const auto &i : adj[cur]) if (i.v != p && !processed[i.v]) {
      mx = max(mx, {sz[i.v], i.v});
    }
    if (mx.first * 2 <= sz[n]) {
      return cur;
    }
    p = cur;
    cur = mx.second;
  }
}

void Solve(int R) {
  function<lint(int, int, vector<lint>&)> dfs = [&](int n, int p, vector<lint> &v) {
    lint cur = 0;
    for (const auto &i : adj[n]) if (i.v != p && !processed[i.v]) {
      lint nxt = dfs(i.v, n, v) + i.away;
      if (cur < nxt) swap(cur, nxt);
      if (nxt > 0) v.emplace_back(nxt);
    }
    return cur;
  };

  vector<lint> all;
  for (const auto &i : adj[R]) if (!processed[i.v]) {
    optimal[i.v] = dfs(i.v, R, all) + i.away;
    all.emplace_back(optimal[i.v]);
  }
  sort(begin(all), end(all), greater<lint>());

  { // Case #1: R is chosen
    lint sum = BaseCost[R]; 
    int cnt = 1;
    Answer[cnt] = max(Answer[cnt], sum);

    for (int i = 0; i < all.size(); i++) {
      sum += all[i], cnt++;
      Answer[cnt] = max(Answer[cnt], sum);
    }
  }

  { // Case #2: R is not chosen, two different subtrees
    pair<lint, lint> mx = {-1, -1};
    for (const auto &i : adj[R]) if (!processed[i.v]) {
      if (optimal[i.v] > mx.first) {
        mx.second = mx.first;
        mx.first = optimal[i.v];
      } else if (optimal[i.v] > mx.second) {
        mx.second = optimal[i.v];
      }
    }
    
    if (mx.first != -1 && mx.second != -1) {
     lint sum = BaseCost[R] + mx.first + mx.second; 
      int cnt = 2;
      Answer[cnt] = max(Answer[cnt], sum);

      for (int i = 0; i < all.size(); i++) {
        if (all[i] == mx.first) {
          mx.first = -1;
        } else if (all[i] == mx.second) {
          mx.second = -1;
        } else {
          sum += all[i], cnt++;
          Answer[cnt] = max(Answer[cnt], sum);
        }
      }
    }
  }
}

void CentroidDecomposition() {
  queue<int> q;
  q.emplace(0);
  while (!q.empty()) {
    int R = Centroid(q.front());
    q.pop();
    Solve(R);
    processed[R] = true;
    for (const auto &i : adj[R]) if (!processed[i.v]) {
      q.emplace(i.v);
    }
  }

  for (int i = 1; i <= N; i++) {
    Answer[i] = max(Answer[i], Answer[i - 1]);
  }
}


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

  cin >> N;

  sz.resize(N);  
  adj.resize(N);
  optimal.resize(N);
  BaseCost.resize(N);
  Answer.resize(N + 1);
  processed.resize(N, false);

  lint sum_all_costs = 0;
  for (int i = 0; i < N - 1; i++) {
    int A, B, C, D;
    cin >> A >> B >> C >> D;
    A--, B--;
    adj[A].emplace_back(B, C, D);
    adj[B].emplace_back(A, D, C);
    sum_all_costs += C + D;
  }

  GenerateBaseCost();
  CentroidDecomposition();

  int Q;
  cin >> Q;
  for (int i = 0; i < Q; i++) {
    int E;
    cin >> E;
    cout << (sum_all_costs - Answer[E]) << "\n";
  }

  return 0;
}

Compilation message (stderr)

designated_cities.cpp: In function 'void Solve(int)':
designated_cities.cpp:95:23: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for (int i = 0; i < all.size(); i++) {
                     ~~^~~~~~~~~~~~
designated_cities.cpp:117:25: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
       for (int i = 0; i < all.size(); i++) {
                       ~~^~~~~~~~~~~~

• Subtask 1: 6 / 6
• Subtask 2: 7 / 7
• Subtask 3: 9 / 9
• Subtask 4: 17 / 17
• Subtask 5: 17 / 17
• Subtask 6: 44 / 44