Submission #1257107

#	Time	Username	Problem	Language	Result	Execution time	Memory
1257107		LucaLucaM	Capital City (JOI20_capital_city)	C++20	11 / 100	2858 ms	30144 KiB

capital_city

#include <iostream>
#include <vector>
#include <algorithm>
#include <cassert>
#include <set>

using ll = long long;
#define debug(x) #x << " = " << x << '\n'

const int INF = 1e9;

int main() {
  std::ios_base::sync_with_stdio(false);
  std::cin.tie(0);

  int n, k;
  std::cin >> n >> k;

  std::vector<std::vector<int>> g(n);

  for (int i = 1; i < n; i++) {
    int u, v;
    std::cin >> u >> v;
    u--, v--;
    g[u].push_back(v);
    g[v].push_back(u);
  }

  std::vector<int> c(n);
  for (int &x : c)  {
    std::cin >> x;
    x--;
  }

  std::vector<bool> blocked(n, false);
  std::vector<int> sz(n);

  auto calc_sizes = [&](auto &&self, int u, int p) -> void {
    sz[u] = 1;
    for (int v : g[u]) {
      if (v != p && !blocked[v]) {
        self(self, v, u);
        sz[u] += sz[v];
      }
    }
  };

  auto find_centroid = [&](auto &&self, int u, int p, int tsz) -> int {
    for (int v : g[u]) {
      if (v != p && !blocked[v] && sz[v] * 2 > tsz) { 
        return self(self, v, u, tsz);
      } 
    }
    return u;
  };  

  std::vector<bool> blockedcolor(k, false);
  std::vector<int> parent(n);

  std::vector<int> vertices;

  auto calc_dfs = [&](auto &&self, int u) -> void {
    vertices.push_back(u);
    for (int v : g[u]) {
      if (v != parent[u] && !blocked[v]) {
        parent[v] = u;
        self(self, v);
      }
    }
  };

  std::vector<bool> vis(n, false);
  std::vector<bool> inSubtree(n, false);

  std::vector<std::vector<int>> thisColor(k);
  for (int i = 0; i < n; i++) {
    thisColor[c[i]].push_back(i);
  }

  std::vector<int> f(k, 0);

  int answer = k - 1;
  auto build_decomp = [&](auto &&self, int u) -> void {
    calc_sizes(calc_sizes, u, -1);
    int centroid = find_centroid(find_centroid, u, -1, sz[u]);
    parent[centroid] = -1;
    calc_dfs(calc_dfs, centroid);

    if (!blockedcolor[c[centroid]]) {
      std::vector<int> st;
      std::vector<int> st2;
      for (int u : vertices) {
        f[c[u]]++;
        inSubtree[u] = true;
      }
      int cur = 0;
      if (f[c[centroid]] < (int) thisColor[c[centroid]].size()) {
        cur = 1e9;
      } else {
        for (int u : thisColor[c[centroid]]) {
          st.push_back(u);
          vis[u] = true;
        }
      }
      while (!st.empty() && cur < answer) {
        int u = st.back();
        st2.push_back(u);
        st.pop_back();
        if (parent[u] != -1 && !vis[parent[u]]) {
          cur++;
          if (f[c[parent[u]]] < (int) thisColor[c[parent[u]]].size()) {
            cur = 1e9;
          } else {
            for (int v : thisColor[c[parent[u]]]) {
              vis[v] = true;
              st.push_back(v);
            }
          }
        }
      }
      for (int x : st) {
        vis[x] = false;
      }
      for (int x : st2) {
        if (!inSubtree[x]) {
          cur = 1e9;
        }
        vis[x] = false;
      }
      answer = std::min(answer, cur);
      for (int x : vertices) {
        inSubtree[x] = false;
        f[c[u]]--;
      }
    }
    
    blocked[centroid] = true;
    blockedcolor[c[centroid]] = true;
    vertices.clear();
    for (int v : g[centroid]) {
      if (!blocked[v]) {
        self(self, v);
      }
    }
  };

  build_decomp(build_decomp, 0);

  std::cout << answer;

  return 0;
}

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