제출 #1300695

#제출 시각아이디문제언어결과실행 시간메모리
1300695AksLolCoding수도 (JOI20_capital_city)C++17
100 / 100
506 ms29460 KiB
#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>

using namespace std;

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

  int n, k;
  cin >> n >> k;
  vector<vector<int>> adj(n + 1);
  for (int i = 0, u, v; i < n - 1; ++i) {
    cin >> u >> v;
    adj[u].push_back(v);
    adj[v].push_back(u);
  }

  vector<int> c(n + 1), count_in(k + 1);
  for (int i = 1; i <= n; ++i) {
    cin >> c[i];
    count_in[c[i]]++;
  }

  vector<bool> is_alr_centroid(n + 1);
  vector<int> subtree_size(n + 1);
  auto find_subtree_size = [&](auto &&self, int u, int p) -> int {
    subtree_size[u] = 1;
    for (int &i : adj[u]) {
      if (i == p || is_alr_centroid[i]) {
        continue;
      }
      subtree_size[u] += self(self, i, u);
    }
    return subtree_size[u];
  };

  auto find_centroid = [&](auto &&self, int n, int u, int p) -> int {
    pair<int, int> best = {0, 0};
    for (int &i : adj[u]) {
      if (i == p || is_alr_centroid[i]) {
        continue;
      }
      best = max(best, {subtree_size[i], i});
    }
    if (best.first <= n / 2) {
      return u;
    }
    return self(self, n, best.second, u);
  };

  int ans = n;
  vector<bool> used_color(k + 1), vis(n + 1);
  vector<int> par(n + 1), cur_in(k + 1);
  vector<vector<int>> towns(k + 1);
  auto check_node = [&](int n, int u) {
    par[u] = u;
    vector<int> colors_seen, nodes_seen;
    auto dfs = [&](auto &&self, int u, int p) -> void {
      colors_seen.push_back(c[u]);
      nodes_seen.push_back(u);
      towns[c[u]].push_back(u);
      for (int &i : adj[u]) {
        if (i == p || is_alr_centroid[i]) {
          continue;
        }
        par[i] = u;
        self(self, i, u);
      }
    };
    dfs(dfs, u, 0);
    vector<int> colors;
    queue<int> q;
    colors.push_back(c[u]);
    used_color[c[u]] = true;
    for (int &i : towns[c[u]]) {
      q.push(i);
    }
    while (!q.empty()) {
      int i = q.front();
      q.pop();
      if (vis[i] || is_alr_centroid[i]) {
        continue;
      }
      cur_in[c[i]]++;
      vis[i] = true;
      if (!used_color[c[i]]) {
        colors.push_back(c[i]);
        used_color[c[i]] = true;
        for (int &city : towns[c[i]]) {
          q.push(city);
        }
      }
      if (!vis[par[i]] && !is_alr_centroid[par[i]]) {
        q.push(par[i]);
      }
    }
    bool good = true;
    for (int &c : colors) {
      good &= cur_in[c] == count_in[c];
      cur_in[c] = 0;
      used_color[c] = false;
    }
    for (int &c : colors_seen) {
      towns[c].clear();
    }
    for (int &u : nodes_seen) {
      vis[u] = false;
    }
    if (good) {
      ans = min(ans, int(colors.size()));
    }
  };

  auto centroid_decomp = [&](auto &&self, int u) -> void {
    find_subtree_size(find_subtree_size, u, 0);
    int centroid = find_centroid(find_centroid, subtree_size[u], u, 0);
    check_node(subtree_size[u], centroid);
    is_alr_centroid[centroid] = true;
    for (int &i : adj[centroid]) {
      if (is_alr_centroid[i]) {
        continue;
      }
      self(self, i);
    }
  };

  centroid_decomp(centroid_decomp, 1);
  cout << ans - 1 << '\n';
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...