Submission #217169

#	Time	Username	Problem	Language	Result	Execution time	Memory
217169		rama_pang	Capital City (JOI20_capital_city)	C++14	100 / 100	898 ms	38524 KiB

capital_city

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

class CapitalCity {
 
 private:

  int N, K;
  vector<int> color;

  vector<vector<int>> adj;
  vector<vector<int>> color_adj;
  
  class Visited {

   private:

    int visited_mark = 1;
    vector<int> visited;

   public:

    Visited() {}

    void Initialize(int n) {
      visited.resize(n);
    }

    void Clear() {
      visited_mark++;
    }

    bool IsIt(int n) {
      return n == -1 || visited[n] == visited_mark;
    }

    void Mark(int n) {
      if (n != -1) visited[n] = visited_mark;
    }

  };

  Visited visited, reachable;

  vector<int> parent;
  vector<int> sz;
  vector<bool> deleted;

  void Dfs(int n, bool r = true) {
    if (r) {
      assert(!deleted[n]);
      parent[n] = -1;
      reachable.Clear();
      reachable.Mark(n);
    }

    sz[n] = 1;
    for (auto &i : adj[n]) if (i != parent[n] && !deleted[i]) {
      reachable.Mark(i);
      parent[i] = n;
      Dfs(i, false);
      sz[n] += sz[i];
    }
  }

  int SolveColor(int root, int k) {
    assert(color[root] == k);
    Dfs(root);

    queue<int> q;
    visited.Clear();

    for (auto &c : color_adj[k]) {
      if (reachable.IsIt(c)) {
        q.emplace(c);
        visited.Mark(c);
      } else {
        return -1;
      }
    }

    int ans = 0;

    while (!q.empty()) {
      int u = q.front();
      q.pop();
      if (!visited.IsIt(parent[u])) {
        ans++;
        for (auto &c : color_adj[color[parent[u]]]) {
          if (reachable.IsIt(c)) {
            q.emplace(c);
            visited.Mark(c);
          } else {
            return -1;
          }
        }
      }
    }

    return ans;
  }

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

  int CentroidDecomposition() {
    set<int> can;
    for (int i = 0; i < N; i++) {
      can.emplace(i);
    }

    auto Delete = [&](int x) {
      for (auto &c : color_adj[x]) {
        can.erase(c);
        deleted[c] = true;
      }
    };
    
    int ans = N - 1;

    while (!can.empty()) {
      int u = *begin(can);
      u = Centroid(u);
      int current = SolveColor(u, color[u]);
      if (current != -1) ans = min(ans, current);
      Delete(color[u]);
    }

    return ans;
  }

  void Init() {
    visited.Initialize(N);
    reachable.Initialize(N);

    parent.resize(N);
    sz.resize(N);
    deleted.resize(N);
  }

 public:

  void Read() {
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);

    cin >> N >> K;

    adj.resize(N);
    color.resize(N);
    color_adj.resize(K);

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

    for (int i = 0; i < N; i++) {
      cin >> color[i];
      color[i]--;
      color_adj[color[i]].emplace_back(i);
    }
  }

  int Solve() {
    Init();
    return CentroidDecomposition();
  }

};

int main() {
  CapitalCity Solver;
  Solver.Read();
  cout << Solver.Solve() << "\n";
  return 0;
}

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 30 / 30
• Subtask 4: 59 / 59