Dynamic connectivity for the Codeforces problem 100551A (Connect and disconnect)

dynamic.cpp

// Dynamic connectivity on a graph
// Solution method: Offline processing + segment tree + undo on union-find
// Details:
//     First, scan through the input to find ranges in which an edge is active
//     Then insert these ranges into the nodes of a segment tree
//     If we merge all the edges in the nodes on a path from root to a leaf,
//         we get the number of connected components at that point
//     We find the component count for all leaves efficiently by adding undo
//         operations to the union-find data structure

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

// Union-find data structure with undo
vector<int> par;
using change = tuple<int, int>;  // (u, old_par)

// Find the root of the tree containing u
// Do path compression and remember the changes to undo later
int find(int u, vector<change>& changes) {
  if (par[u] == u) return u;
  int v = find(par[u], changes);
  if (par[u] == v) return v;

  // Path compression should happen here: So remember to undo later
  changes.push_back({u, par[u]});
  return par[u] = v;
}

// Merge components of u and v, and remember the changes to undo later
// Return 1 if a merge happened, 0 otherwise
int merge(int u, int v, vector<change>& changes) {
  u = find(u, changes);
  v = find(v, changes);
  if (u == v) return 0;

  // Merge changes the parent, so remember to undo later
  changes.push_back({u, u});
  par[u] = v;
  return 1;
}

// Segment tree to store the active edges in a range
using edge = pair<int, int>;
int base;
vector<vector<edge>> seg;

// Insert edge (u, v) in the segment tree node [l, r]
// After the insert operation, edge will be present in every node whose range
//     is completely contained in the range of the edge
void insert(int node, int l, int r, int L, int R, int u, int v) {
  if (l == L and r == R) {
    seg[node].push_back({u, v});
    return;
  }
  int M = (L + R) >> 1, lc = node << 1, rc = lc | 1;
  if (r <= M) {
    insert(lc, l, r, L, M, u, v);
  } else if (l > M) {
    insert(rc, l, r, M + 1, R, u, v);
  } else {
    insert(lc, l, M, L, M, u, v);
    insert(rc, M + 1, r, M + 1, R, u, v);
  }
}

// Main logic: Process the path from the root to leaves of the segment tree
// Merging will reduce the component count, then recurse into children
// Undo the changes after processing the children to return to parent state
void process(int node, vector<int>& answer, int components) {
  // Merge and reduce component count, and remember the changes
  vector<change> changes;
  for (auto [u, v] : seg[node]) components -= merge(u, v, changes);

  if (node >= base) {
    // If we are at a leaf node, store the component count
    answer[node - base] = components;
  } else {
    // Otherwise, recurse into children
    int lc = node << 1, rc = lc | 1;
    process(lc, answer, components);
    process(rc, answer, components);
  }

  // Undo the DSU chnages in reverse order to return to original state
  while (!changes.empty()) {
    auto [u, a] = changes.back();
    changes.pop_back();
    par[u] = a;
  }
}

int main() {
  ifstream fin("connect.in");
  ofstream fout("connect.out");

  long long N, K;
  fin >> N >> K;

  // Initialize the segment tree
  base = 1;
  while (base < K) base <<= 1;
  seg = vector<vector<edge>>(base << 1);

  // Initialize the union-find data structure
  par.resize(N);
  iota(par.begin(), par.end(), 0);

  unordered_map<long long, int> start;  // Insertion time of each edge
  vector<int> queries;                  // Indices of queries

  for (int i = 0; i < K; i++) {
    string c;
    fin >> c;
    if (c == "?") {
      queries.push_back(i);
      continue;
    }
    int u, v;
    fin >> u >> v;
    --u, --v;
    if (u > v) swap(u, v);
    long long e = u * N + v;  // long long to put in hashmap

    if (c == "+") {
      start[e] = i;
    } else {
      // Insert the (u, v) edge at the range [start[e], i - 1]
      insert(1, start[e], i - 1, 0, base - 1, u, v);
      start.erase(e);
    }
  }

  // Insert the remaining edges that are active at the end
  for (auto [e, i] : start) insert(1, i, K - 1, 0, base - 1, e / N, e % N);

  // Process the segment tree to find the component count after each query
  vector<int> answer(base);
  process(1, answer, N);

  // Output the answer for '?' queries
  for (int& q : queries) fout << answer[q] << '\n';
  return 0;
}