Submission #219880

#	Time	Username	Problem	Language	Result	Execution time	Memory
219880		rama_pang	Mergers (JOI19_mergers)	C++14	10 / 100	3081 ms	14056 KiB

mergers

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

class DisjointSet {
 private:
  int n;
  vector<int> p;

 public:
  DisjointSet(int n) : n(n) {
    p.resize(n);
    iota(begin(p), end(p), 0);
  }

  int Find(int x) {
    return p[x] == x ? x : p[x] = Find(p[x]);
  }

  bool Union(int x, int y) {
    x = Find(x);
    y = Find(y);
    if (x == y) return false;
    p[x] = y;
    return true;
  }

  bool IsRoot(int x) {
    return p[x] == x;
  }
};

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

  int N, K;
  cin >> N >> K;

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

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

  vector<int> sz(N, 0);
  vector<int> par(N, -1);
  vector<bool> visColor(K, false);
  DisjointSet group(N);

  function<void(int)> Dfs = [&](int n) {
    sz[n] = 1;
    for (auto i : adj[n]) if (i != par[n] && !visColor[color[i]]) {
      par[i] = n;
      Dfs(i);
      sz[n] += sz[i];
    }
  };

  auto Centroid = [&](int r) {
    par[r] = -1;
    Dfs(r);
    int Size = sz[r];

    while (true) {
      pair<int, int> mx = {-1, -1};
      for (auto i : adj[r]) if (i != par[r] && !visColor[color[i]]) {
        mx = max(mx, {sz[i], i});
      }
      if (mx.first * 2 <= Size) {
        return r;
      }
      r = mx.second;
    }
  };

  auto Solve = [&](int r) {
    par[r] = -1;
    Dfs(r);

    queue<int> q;

    visColor[color[r]] = true;
    for (auto i : colorList[color[r]]) {
      q.emplace(i);
      group.Union(i, r);
    }

    while (!q.empty()) {
      int p = par[q.front()];
      q.pop();

      if (p != -1 && !visColor[color[p]]) {
        visColor[color[p]] = true;
        for (auto i : colorList[color[p]]) {
          q.emplace(i);
          group.Union(r, i);
        }
      } else if (p != -1 && visColor[color[p]]) {
        group.Union(r, p); 
      }
    }
  };

  auto CentroidDecomposition = [&]() {
    for (int i = 0; i < N; i++) {
      while (!visColor[color[i]]) {
        Solve(i);
      }
    }
  };

  CentroidDecomposition();

  vector<int> degree(N, 0);
  for (int i = 0; i < N - 1; i++) {
    int u, v;
    tie(u, v) = edges[i];
    if (group.Find(u) != group.Find(v)) {
      degree[group.Find(u)]++;
      degree[group.Find(v)]++;
    }
  }

  int ans = 0;
  for (int i = 0; i < N; i++) {
    if (group.IsRoot(i) && degree[i] == 1) {
      ans++;
    }
  }

  cout << ((ans + 1) / 2) << "\n";
  return 0;
}

Compilation message (stderr)

mergers.cpp: In function 'int main()':
mergers.cpp:71:8: warning: variable 'Centroid' set but not used [-Wunused-but-set-variable]
   auto Centroid = [&](int r) {
        ^~~~~~~~

• Subtask 1: 10 / 10
• Subtask 2: 0 / 24
• Subtask 3: 0 / 14
• Subtask 4: 0 / 22
• Subtask 5: 0 / 30