AIOJ

Mo's Algorithm on Trees

Ordinary Mo's algorithm can handle only linear problems, so we force the tree into a sequence.

We can traverse the tree in parenthesis order, divide that order into blocks, and run Mo's algorithm on it.

How exactly is this done?

Process

Run DFS on a tree. When DFS reaches node x, append push_back(x); when DFS finishes node x, append push_back(-x). When moving the pointers,

when the newly added value is x --->add(x)
when the newly added value is -x --->del(x)
when the newly removed value is x --->del(x)
when the newly removed value is -x --->add(x)

In this way, we convert a tree into a sequence.

Example Problems

Example: 「WC2013」Candy Park

Problem: You are given a tree whose $i$ -th vertex has color $c_i$ . For each query on a path $u_i$ , $v_i$ , find

$\sum_{c}val_c\sum_{i=1}^{cnt_c}w_i$

Here, $val$ is the value of the color, $cnt$ is the number of occurrences of the color, and $w$ is the value after the color has appeared $i$ times.

Process

First convert the tree into a sequence. Then, whenever a vertex is added or removed, its contribution to the answer can be obtained in $O(1)$ time, namely $val_c\times w_{cnt_{c+1}}$ .

We notice that the traversal also scans the subtree of the starting point, producing extra contributions. What should we do?

During the scan, vertices in the starting point's subtree are certainly scanned twice, but their contribution is zero.

Therefore, we can create a $vis$ array and XOR $vis_x$ with 1 every time vertex x is scanned.

If $vis_x=0$ , the contribution of this vertex can be ignored.

Thus, we can solve the problem with Mo's algorithm on trees.

For modifications, add a time dimension; this gives Mo's algorithm on trees with modifications.

The interval may not contain the LCA; in that case, remove the extra contribution, and that is all.

Implementation

Reference Code

#include <algorithm>
#include <cmath>
#include <cstdio>
using namespace std;

constexpr int MAXN = 200010;

int f[MAXN], g[MAXN], id[MAXN], head[MAXN], cnt, last[MAXN], dep[MAXN],
    fa[MAXN][22], v[MAXN], w[MAXN];
int block, index, n, m, q;
int pos[MAXN], col[MAXN], app[MAXN];
bool vis[MAXN];
long long ans[MAXN], cur;

struct edge {
  int to, nxt;
} e[MAXN];

int cnt1 = 0, cnt2 = 0;  // Timestamps

struct query {
  int l, r, t, id;

  bool operator<(const query &b) const {
    return (pos[l] < pos[b.l]) || (pos[l] == pos[b.l] && pos[r] < pos[b.r]) ||
           (pos[l] == pos[b.l] && pos[r] == pos[b.r] && t < b.t);
  }
} a[MAXN], b[MAXN];

void addedge(int x, int y) {
  e[++cnt] = edge{y, head[x]};
  head[x] = cnt;
}

void dfs(int x) {
  id[f[x] = ++index] = x;
  for (int i = head[x]; i; i = e[i].nxt) {
    if (e[i].to != fa[x][0]) {
      fa[e[i].to][0] = x;
      dep[e[i].to] = dep[x] + 1;
      dfs(e[i].to);
    }
  }
  id[g[x] = ++index] = x;  // Parenthesis order
}

int lca(int x, int y) {
  if (dep[x] < dep[y]) swap(x, y);
  if (dep[x] != dep[y]) {  // Raise to the same depth
    int dis = dep[x] - dep[y];
    for (int i = 20; i >= 0; i--)
      if (dis >= (1 << i)) dis -= 1 << i, x = fa[x][i];
  }
  if (x == y) return x;
  for (int i = 20; i >= 0; i--) {
    if (fa[x][i] != fa[y][i]) x = fa[x][i], y = fa[y][i];
  }
  return fa[x][0];
}

void add(int x) {
  if (vis[x])
    cur -= (long long)v[col[x]] * w[app[col[x]]--];
  else
    cur += (long long)v[col[x]] * w[++app[col[x]]];
  vis[x] ^= 1;
}

// Move along the time dimension
void modify(int x, int t) {
  if (vis[x]) {
    add(x);
    col[x] = t;
    add(x);
  } else
    col[x] = t;
}

int main() {
  scanf("%d%d%d", &n, &m, &q);
  for (int i = 1; i <= m; i++) scanf("%d", &v[i]);
  for (int i = 1; i <= n; i++) scanf("%d", &w[i]);
  for (int i = 1; i < n; i++) {
    int x, y;
    scanf("%d%d", &x, &y);
    addedge(x, y);
    addedge(y, x);
  }
  for (int i = 1; i <= n; i++) {
    scanf("%d", &last[i]);
    col[i] = last[i];
  }
  dfs(1);
  for (int j = 1; j <= 20; j++)
    for (int i = 1; i <= n; i++)
      fa[i][j] = fa[fa[i][j - 1]][j - 1];  // Precompute ancestors
  int block = pow(index, 2.0 / 3);
  for (int i = 1; i <= index; i++) {
    pos[i] = (i - 1) / block;
  }
  while (q--) {
    int opt, x, y;
    scanf("%d%d%d", &opt, &x, &y);
    if (opt == 0) {
      b[++cnt2].l = x;
      b[cnt2].r = last[x];
      last[x] = b[cnt2].t = y;
    } else {
      if (f[x] > f[y]) swap(x, y);
      a[++cnt1] = query{lca(x, y) == x ? f[x] : g[x], f[y], cnt2, cnt1};
    }
  }
  sort(a + 1, a + cnt1 + 1);
  int L, R, T;  // Pointer coordinates
  L = R = 0;
  T = 1;
  for (int i = 1; i <= cnt1; i++) {
    while (T <= a[i].t) {
      modify(b[T].l, b[T].t);
      T++;
    }
    while (T > a[i].t) {
      modify(b[T].l, b[T].r);
      T--;
    }
    while (L > a[i].l) {
      L--;
      add(id[L]);
    }
    while (L < a[i].l) {
      add(id[L]);
      L++;
    }
    while (R > a[i].r) {
      add(id[R]);
      R--;
    }
    while (R < a[i].r) {
      R++;
      add(id[R]);
    }
    int x = id[L], y = id[R];
    int llca = lca(x, y);
    if (x != llca && y != llca) {
      add(llca);
      ans[a[i].id] = cur;
      add(llca);
    } else
      ans[a[i].id] = cur;
  }
  for (int i = 1; i <= cnt1; i++) {
    printf("%lld\n", ans[i]);
  }
  return 0;
}

Genuine Mo's Algorithm on Trees

The tree Mo's algorithm above only converts the tree into a chain; the following is the genuine tree version.

Because Mo-related problems are template problems, the implementation is not explained in much detail.

Query Ordering

We know that Mo's algorithm is based on block decomposition, so we need to find a decomposition of the tree that guarantees the time complexity.

Conditions:

The distance between nodes in the same block must not exceed the given block size.
Each block must contain neither too many nor too few nodes.
Every node must belong to a block.
The distance between consecutively numbered blocks must not be too large.

With these conditions in mind, consider 「SCOI2005」Kingdom Federation.

Based on this problem, we only need to ensure the last condition to solve the block-decomposition problem.

Idea

Let lim be the desired block size. First, run DFS over the entire tree; whenever a subtree's size exceeds lim, put its nodes into one block. It is easy to see that some nodes may remain at the root, so put those nodes into the last block.

Method: use a stack to maintain the current node while visiting its children. When the distance from the top of the stack to the parent exceeds the desired block size, pop these elements and make them a block; make the remaining elements a block of their own at the end.

The final sorting method is as follows: if the timestamp in the first dimension is greater than that in the second, swap them; sort by the block containing the first dimension as the primary key and the timestamp in the second dimension as the secondary key.

Moving the Pointers

Process

It is natural to mark the vertices counted in the answer, let the pointers move directly toward their targets, and toggle the vertices on the paths.

However, there is a problem. If both pointers start at x, then x is clearly marked. When the two pointers move to the same child (among many other cases), x should not be marked, but in fact x is marked because the two pointers each mark it once, causing the two toggles to cancel.

How can we solve this?

There is an obvious property: these vertices must be some LCAs, because only at an LCA can repeated unmarking cause the unmarking to fail.

Therefore, we do not mark the LCA each time; instead, mark it when the answer needs to be queried, and then unmark it afterward.

Implementation

// Toggle every node on the path except the LCA
void move(int x, int y) {
  if (dp[x] < dp[y]) swap(x, y);
  while (dp[x] > dp[y]) update(x), x = fa[x];
  while (x != y) update(x), update(y), x = fa[x], y = fa[y];
  // x != y guarantees that the LCA is not toggled
}

For computing the LCA, we can use heavy-light decomposition. Then we can place the block-decomposition steps in the first DFS of the heavy-light decomposition, and use the DFS order from the second DFS directly as the timestamps.

int bl[100002], bls = 0;  // The block containing the node and the number of blocks
unsigned step;            // Block size
int fa[100002], dp[100002], hs[100002] = {0}, sz[100002] = {0};
// Parent, depth, heavy child, and size
stack<int> sta;

void dfs1(int x) {
  sz[x] = 1;
  unsigned ss = sta.size();
  for (int i = head[x]; i; i = nxt[i])
    if (ver[i] != fa[x]) {
      fa[ver[i]] = x;
      dp[ver[i]] = dp[x] + 1;
      dfs1(ver[i]);
      sz[x] += sz[ver[i]];
      if (sz[ver[i]] > sz[hs[x]]) hs[x] = ver[i];
      if (sta.size() - ss >= step) {
        bls++;
        while (sta.size() != ss) bl[sta.top()] = bls, sta.pop();
      }
    }
  sta.push(x);
}

// main
if (!sta.empty()) {
  bls++;  // This line is optional
  while (!sta.empty()) bl[sta.top()] = bls, sta.pop();
}

Time Complexity

Now we reach the key point, which concerns the choice of block size.

Let the block size be $unit$ :

For the x pointer, because the distances between nodes in each block are around $unit$ , the x pointer moves $unit^2$ times in each block ( $unit\times dis_{\max}$ ), for a total of $n\times unit$ moves ( $unit^2 \times (\frac{n}{unit})$ );
The y pointer moves at most $O(n)$ times in each block, for a total of $\frac{n^2}{unit}$ moves ( $n \times (\frac{n}{unit})$ ).

Their sum is minimized approximately at the square-root point (because the block sizes in tree Mo's algorithm are not fixed, it need not be followed exactly).

Example: 「WC2013」Candy Park

Because there is an additional time dimension, a block size around $n^{0.6}$ is appropriate.

Reference Code

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <stack>
using namespace std;

int gi() {
  int x, c, op = 1;
  while (c = getchar(), c < '0' || c > '9')
    if (c == '-') op = -op;
  x = c ^ 48;
  while (c = getchar(), c >= '0' && c <= '9')
    x = (x << 3) + (x << 1) + (c ^ 48);
  return x * op;
}

int head[100002], nxt[200004], ver[200004], tot = 0;

void add(int x, int y) {
  ver[++tot] = y, nxt[tot] = head[x], head[x] = tot;
  ver[++tot] = x, nxt[tot] = head[y], head[y] = tot;
}

int bl[100002], bls = 0;
unsigned step;
int fa[100002], dp[100002], hs[100002] = {0}, sz[100002] = {0}, top[100002],
                            id[100002];
stack<int> sta;

void dfs1(int x) {
  sz[x] = 1;
  unsigned ss = sta.size();
  for (int i = head[x]; i; i = nxt[i])
    if (ver[i] != fa[x]) {
      fa[ver[i]] = x, dp[ver[i]] = dp[x] + 1;
      dfs1(ver[i]);
      sz[x] += sz[ver[i]];
      if (sz[ver[i]] > sz[hs[x]]) hs[x] = ver[i];
      if (sta.size() - ss >= step) {
        bls++;
        while (sta.size() != ss) bl[sta.top()] = bls, sta.pop();
      }
    }
  sta.push(x);
}

int cnt = 0;

void dfs2(int x, int hf) {
  top[x] = hf, id[x] = ++cnt;
  if (!hs[x]) return;
  dfs2(hs[x], hf);
  for (int i = head[x]; i; i = nxt[i])
    if (ver[i] != fa[x] && ver[i] != hs[x]) dfs2(ver[i], ver[i]);
}

int lca(int x, int y) {
  while (top[x] != top[y]) {
    if (dp[top[x]] < dp[top[y]]) swap(x, y);
    x = fa[top[x]];
  }
  return dp[x] < dp[y] ? x : y;
}

struct qu {
  int x, y, t, id;

  bool operator<(const qu a) const {
    return bl[x] == bl[a.x] ? (bl[y] == bl[a.y] ? t < a.t : bl[y] < bl[a.y])
                            : bl[x] < bl[a.x];
  }
} q[100001];

int qs = 0;

struct ch {
  int x, y, b;
} upd[100001];

int ups = 0;
long long ans[100001];
int b[100001] = {0};
int a[100001];
long long w[100001];
long long v[100001];
long long now = 0;
bool vis[100001] = {false};

void back(int t) {
  if (vis[upd[t].x]) {
    now -= w[b[upd[t].y]--] * v[upd[t].y];
    now += w[++b[upd[t].b]] * v[upd[t].b];
  }
  a[upd[t].x] = upd[t].b;
}

void change(int t) {
  if (vis[upd[t].x]) {
    now -= w[b[upd[t].b]--] * v[upd[t].b];
    now += w[++b[upd[t].y]] * v[upd[t].y];
  }
  a[upd[t].x] = upd[t].y;
}

void update(int x) {
  if (vis[x])
    now -= w[b[a[x]]--] * v[a[x]];
  else
    now += w[++b[a[x]]] * v[a[x]];
  vis[x] ^= 1;
}

void move(int x, int y) {
  if (dp[x] < dp[y]) swap(x, y);
  while (dp[x] > dp[y]) update(x), x = fa[x];
  while (x != y) update(x), update(y), x = fa[x], y = fa[y];
}

int main() {
  int n = gi(), m = gi(), k = gi();
  step = (int)pow(n, 0.6);
  for (int i = 1; i <= m; i++) v[i] = gi();
  for (int i = 1; i <= n; i++) w[i] = gi();
  for (int i = 1; i < n; i++) add(gi(), gi());
  for (int i = 1; i <= n; i++) a[i] = gi();
  for (int i = 1; i <= k; i++)
    if (gi())
      q[++qs].x = gi(), q[qs].y = gi(), q[qs].t = ups, q[qs].id = qs;
    else
      upd[++ups].x = gi(), upd[ups].y = gi();
  for (int i = 1; i <= ups; i++) upd[i].b = a[upd[i].x], a[upd[i].x] = upd[i].y;
  for (int i = ups; i; i--) back(i);
  fa[1] = 1;
  dfs1(1), dfs2(1, 1);
  if (!sta.empty()) {
    bls++;
    while (!sta.empty()) bl[sta.top()] = bls, sta.pop();
  }
  for (int i = 1; i <= n; i++)
    if (id[q[i].x] > id[q[i].y]) swap(q[i].x, q[i].y);
  sort(q + 1, q + qs + 1);
  int x = 1, y = 1, t = 0;
  for (int i = 1; i <= qs; i++) {
    if (x != q[i].x) move(x, q[i].x), x = q[i].x;
    if (y != q[i].y) move(y, q[i].y), y = q[i].y;
    int f = lca(x, y);
    update(f);
    while (t < q[i].t) change(++t);
    while (t > q[i].t) back(t--);
    ans[q[i].id] = now;
    update(f);
  }
  for (int i = 1; i <= qs; i++) printf("%lld\n", ans[i]);
  return 0;
}