Submission #1257089

#	Time	Username	Problem	Language	Result	Execution time	Memory
1257089		LucaLucaM	수도 (JOI20_capital_city)	C++20	0 / 100	3096 ms	30876 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;

namespace DSU {
  std::vector<int> p;
  std::vector<int> sz;
  std::vector<std::set<int>> st;
  std::vector<int> need;
  std::vector<int> f;

  void init(int n, const std::vector<int> &colors) {
    p.resize(n);
    f.assign(n, 0);
    sz.assign(n, 1);
    st.resize(n);
    need.resize(n);
    for (int x : colors) {
      f[x]++;
    }
    for (int i = 0; i < n; i++) {
      p[i] = i;
      st[i].insert(colors[i]);
      need[i] = f[colors[i]];
    }
  }

  int root(int u) {
    return p[u] == u? u : p[u] = root(p[u]);
  }
  void join(int u, int v) {
    u = root(u);
    v = root(v);
    if (u == v) {
      return;
    }
    if (st[u].size() < st[v].size()) {
      std::swap(u, v);
    }
    for (int c : st[v]) {
      if (!st[u].count(c)) {
        st[u].insert(c);
        need[u] += f[c];
      }
    }
    sz[u] += sz[v];
    p[v] = u;
  }
  bool isGood(int u) {
    u = root(u);
    return need[u] == sz[u];
  }
  int countDistinct(int u) {
    u = root(u);
    return (int) st[u].size();
  }
};

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);
  }

  int answer = n;
  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) {
        inSubtree[u] = true;
      }
      int need = (int) thisColor[c[centroid]].size();
      int cur = 0;
      for (int u : thisColor[c[centroid]]) {
        if (centroid == 0) {
          std::cout << debug(u);
        }
        if (!inSubtree[u]) {
          cur = 1e9;
          break;
        }
        st.push_back(u);
        vis[u] = true;
      }
      while (!st.empty()) {
        int u = st.back();
        st2.push_back(u);
        st.pop_back();
        if (parent[u] != -1 && !vis[parent[u]]) {
          if (blockedcolor[c[parent[u]]]) {
            cur = 1e9;
            break;
          }
          cur++;
          need += (int) thisColor[c[parent[u]]].size();
          for (int v : thisColor[c[parent[u]]]) {
            if (!inSubtree[v]) {
              cur = 1e9;
              break;
            }
            vis[v] = true;
            st.push_back(v);
          }
        }
      }
      for (int x : st2) {
        vis[x] = false;
      }
      // if (cur == 0 && centroid == 0) {
        // std::cout << "yesSir\n";
        // for (int x : st2) {
        //   std::cout << debug(x);
        // }
      // }
      if (need == (int) st2.size()) {
        // if (cur == 0) {
          // std::cout << " ???\n";
          // for (int x : st2) {
          //   std::cout << debug(x);
          // }
        // }
        answer = std::min(answer, cur);
      }
      for (int u : vertices) {
        inSubtree[u] = false;
      }
    }

    // std::cout << debug(centroid);
    
    blocked[centroid] = true;
    blockedcolor[c[centroid]] = true;
    for (int v : g[centroid]) {
      if (!blocked[v]) {
        self(self, v);
      }
    }
  };

  build_decomp(build_decomp, 0);

  std::cout << answer;
  return 0;

  // std::reverse(order.begin(), order.end());

  // DSU::init(n, c);

  // int answer = k - 1;
  // for (int i = 0; i < n; i++) {
  //   if (DSU::isGood(i)) {
  //     answer = 0;
  //     std::cout << 0;
  //     return 0;
  //   }
  // }

  // std::vector<bool> open(n, false);
  
  // for (int u : order) {
  //   std::cout << debug(u);
  //   for (int v : thisColor[c[u]]) {
  //     open[v] = true;
  //     for (int w : g[v]) {
  //       if (open[w]) {
  //         DSU::join(v, w);
  //       }
  //     }
  //   }
  //   if (DSU::isGood(u))  {
  //     answer = std::min(answer, DSU::countDistinct(u) - 1);
  //   }
  // }

  // std::cout << answer;
  
  return 0;
}

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