Submission #1194634

#	Time	Username	Problem	Language	Result	Execution time	Memory
1194634		avighna	Mergers (JOI19_mergers)	C++20	100 / 100	741 ms	196908 KiB

mergers

#include <bits/stdc++.h>

class UFDS {
public:
  std::vector<int> par;
  int n;

  UFDS(int n) : n(n + 1), par(n + 1) { std::iota(par.begin(), par.end(), 0); }

  int root(int u) { return u == par[u] ? u : par[u] = root(par[u]); }

  void merge(int u, int v) {
    u = root(u), v = root(v);
    par[v] = u == v ? par[v] : u;
  }
};

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

  int n, k;
  std::cin >> n >> k;
  std::vector<std::vector<int>> adj(n + 1), nodes(k + 1);
  for (int i = 0, a, b; i < n - 1; ++i) {
    std::cin >> a >> b;
    adj[a].push_back(b), adj[b].push_back(a);
  }
  for (int i = 1, x; i <= n; ++i) {
    std::cin >> x;
    nodes[x].push_back(i);
  }

  std::vector<int> d(n + 1), top(n + 1);
  std::vector up(n + 1, std::vector<int>(20));
  up[1][0] = 1;
  auto dfs = [&](auto &&self, int u, int p) -> void {
    for (auto &i : adj[u]) {
      if (i == p) {
        continue;
      }
      d[i] = d[u] + 1, up[i][0] = u;
      self(self, i, u);
    }
  };
  dfs(dfs, 1, 0);
  for (int bt = 1; bt < 20; ++bt) {
    for (int i = 1; i <= n; ++i) {
      up[i][bt] = up[up[i][bt - 1]][bt - 1];
    }
  }

  auto lca = [&](int a, int b) {
    if (d[a] < d[b]) {
      std::swap(a, b);
    }
    for (int bt = 0, k = d[a] - d[b]; bt < 20; ++bt, k >>= 1) {
      a = k & 1 ? up[a][bt] : a;
    }
    if (a == b) {
      return a;
    }
    for (int bt = 19; bt >= 0; --bt) {
      if (up[a][bt] != up[b][bt]) {
        a = up[a][bt], b = up[b][bt];
      }
    }
    return up[a][0];
  };

  std::priority_queue<std::pair<int, int>> pq; // (depth, node)
  std::iota(top.begin(), top.end(), 0);
  auto add = [&](int u, int p) {
    // go from u -> ... -> p
    if (d[top[u]] <= d[p]) {
      return;
    }
    if (top[u] == u) {
      pq.push({d[u], u});
    }
    top[u] = p;
  };
  for (int i = 1; i <= k; ++i) {
    int l = std::accumulate(nodes[i].begin() + 1, nodes[i].end(), nodes[i][0], lca);
    std::ranges::for_each(nodes[i], [&](int x) { add(x, l); });
  }

  UFDS dsu(n);
  while (!pq.empty()) {
    auto [d, u] = pq.top();
    pq.pop();
    dsu.merge(u, up[u][0]);
    add(up[u][0], top[u]);
  }

  std::vector<std::set<int>> g(n + 1);
  for (int i = 1; i <= n; ++i) {
    for (auto &u : adj[i]) {
      int x = dsu.root(i), y = dsu.root(u);
      if (x != y) {
        g[x].insert(y), g[y].insert(x);
      }
    }
  }
  std::cout << (std::ranges::count_if(g, [](const auto& i) {
    return i.size() == 1;
  }) + 1) / 2 << '\n';
}

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