Submission #1194447

#	Time	Username	Problem	Language	Result	Execution time	Memory
1194447		avighna	Mergers (JOI19_mergers)	C++20	0 / 100	105 ms	38592 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);
    if (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);
  for (int i = 0, a, b; i < n - 1; ++i) {
    std::cin >> a >> b;
    adj[a].push_back(b);
    adj[b].push_back(a);
  }
  std::vector<std::vector<int>> nodes(k + 1);
  for (int i = 0, x; i < n; ++i) {
    std::cin >> x;
    nodes[x].push_back(i + 1);
  }

  std::vector<int> depth(n + 1);
  std::vector lift(n + 1, std::vector<int>(20));
  lift[1][0] = 1;
  auto dfs = [&](auto &&self, int u, int p) -> void {
    for (auto &i : adj[u]) {
      if (i == p) {
        continue;
      }
      depth[i] = depth[u] + 1;
      lift[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) {
      lift[i][bt] = lift[lift[i][bt - 1]][bt - 1];
    }
  }

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

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

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

  std::vector<std::set<int>> adj_new(n + 1);
  for (int i = 1; i <= n; ++i) {
    for (auto &u : adj[i]) {
      int x = dsu.root(i);
      int y = dsu.root(u);
      if (x != y) {
        adj_new[x].insert(y);
        adj_new[y].insert(x);
      }
    }
  }
  int leaves = 0;
  for (int i = 1; i <= n; ++i) {
    leaves += adj_new[i].size() == 1;
  }
  std::cout << std::max(0, leaves - 1) << '\n';
}

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