CSES 1160 Planets Queries II

Published on

Problem Statement

There are $n$ nodes, and you are given an array $t$ of length $n$, where $t_i$ denotes that when you are at node $i$, you can teleport to node $t_i$.

You have to answer $q$ queries. In each query, you are given two integers $u$ and $v$. You have to count how many teleportations it takes to get from $u$ to $v$, or report that it is impossible to reach $v$ from $u$.

$1 \leq n, q \leq 2 \times 10^5$

Tutorial

Obviously, if we connect each node $i$ to node $t_i$, the graph forms a unicyclic forest.

Assume $u$ and $v$ belong to the same connected component.

Now we have four cases:

  1. $u$ and $v$ are both on the cycle
    In this case, $u$ can always reach $v$. We simply need to compute the distance from $u$ to $v$ on the cycle.
  2. $u$ is in the tree and $v$ is on the cycle
    In this case, $u$ can also reach $v$. We need to compute the distance from $u$ to its ancestor on the cycle, then proceed along the cycle to $v$.
  3. $u$ is on the cycle and $v$ is in the tree
    In this case, $u$ can never reach $v$.
  4. $u$ and $v$ are both in the tree
    In this case, we only need to check whether $v$ is an ancestor of $u$.

Time complexity: $\tilde O\left(n\right)$.

Code

struct VOIDENGINE {
  VOIDENGINE([[maybe_unused]] int TEST_NUMBER) {
    int n, m; bin(n, m);
    vi a(n); bin(a);
    vector<vi> e(n);
    FDSU d(n);
    eu(i, n) --a[i];
    eu(i, n) e[a[i]].push_back(i);
    eu(i, n) d.link(i, a[i]);
    vi deg(n);
    eu(i, n) ++deg[a[i]];
    list<int> q;
    eu(i, n) {
      if (deg[i] == 0) q.push_back(i);
    }
    vb cyc(n, true);
    while (q.size() != 0) {
      int x = q.front();
      cyc[x] = false;
      q.pop_front();
      if (--deg[a[x]] == 0) q.push_back(a[x]);
    }
    vi cn;
    eu(i, n) {
      if (cyc[i]) cn.push_back(i);
    }
    vi lid(n, -1);
    vi sz(n);
    auto dfs = lexp([&](auto $, int x, int ff) -> void {
      if (!cyc[x]) return;
      if (lid[x] != -1) return;
      ++sz[d.root(x)];
      lid[x] = lid[ff] + 1;
      $(a[x], x);
    });
    for (auto i : cn) {
      if (lid[i] == -1) {
        lid[i] = 0;
        sz[d.root(i)] = 1;
        if (a[i] != i) dfs(a[i], i);
      }
    }
    vi dep(n);
    vi tl(n);
    auto gdep = lexp([&](auto $, int x, int o) -> void {
      if (x != o && cyc[x]) return;
      tl[x] = o;
      for (auto i : e[x]) {
        if (cyc[i]) continue;
        dep[i] = dep[x] + 1;
        $(i, o);
      }
    });
    for (auto i : cn) gdep(i, i);
    auto jmp = vcc(0, 20, n);
    eu(i, n) jmp[0][i] = a[i];
    for (auto i : cn) jmp[0][i] = i;
    eu(i, 1, 20) {
      eu(j, n) jmp[i][j] = jmp[i - 1][jmp[i - 1][j]];
    }
    eu(m) {
      int u, v; bin(u, v);
      --u; --v;
      if (d.root(u) != d.root(v)) {
        cout << "-1\n";
        continue;
      }
      if (cyc[u] && !cyc[v]) {
        cout << "-1\n";
        continue;
      }
      if (cyc[u] && cyc[v]) {
        int x = lid[u];
        int y = lid[v];
        int z = y - x;
        if (z < 0) z += sz[d.root(v)];
        cout << z << "\n";
        continue;
      }
      if (!cyc[u] && cyc[v]) {
        int x = lid[tl[u]];
        int y = lid[v];
        int z = y - x;
        if (z < 0) z += sz[d.root(v)];
        z += dep[u];
        cout << z << "\n";
        continue;
      }
      if (dep[u] < dep[v]) {
        cout << "-1\n";
        continue;
      }
      int tar = dep[u] - dep[v];
      for (int i = 19; i >= 0; --i) {
        if (dep[jmp[i][u]] >= dep[v]) u = jmp[i][u];
      }
      if (u == v) cout << tar << "\n";
      else cout << "-1\n";
    }
  }
};

void VOIDSPECTRE([[maybe_unused]] int argc, [[maybe_unused]] char* argv[]) {
}

No comments

Post a comment