重鏈剖分題目選講

Otue發表於2024-05-01

染色

給定一棵 \(n\) 個節點的無根樹，共有 \(m\) 個操作，操作分為兩種：

  1. 將節點 \(a\) 到節點 \(b\) 的路徑上的所有點（包括 \(a\) 和 \(b\)）都染成顏色 \(c\)。
  2. 詢問節點 \(a\) 到節點 \(b\) 的路徑上的顏色段數量。

顏色段的定義是極長的連續相同顏色被認為是一段。例如 112221 由三段組成：11、222、1。

這道題是樹剖好題。我們來尋思如何改進樹剖完成題目操作。下面是一個普通樹剖，求路徑和。

int query_uv(int u, int v) {
    int res = 0;
    while (top[u] != top[v]) {
        if (dep[top[u]] < dep[top[v]]) swap(u, v);
        res += query(1, id[top[u]], id[u]);
        u = f[top[u]];
    }
    if (dep[u] < dep[v]) swap(u, v);
    res += query(1, id[v], id[u]);
    return res;
}

我們肯定要將第 5 行線段樹的 query 改成求連續段數量的函式。這個用線段樹維護其實並不困難。見如下程式碼：

struct node {
	int lc, rc, sum;
};
struct segment {
	#define ls p << 1
	#define rs p << 1 | 1
	struct edge {
		int l, r, lc, rc, sum, lazy; // lc左端點顏色，rc右端點顏色，sum顏色端數量
	}tree[N * 4];
	
	void down(int p, int x) {
		tree[p].lc = tree[p].rc = tree[p].lazy = x;
		tree[p].sum = 1;
	}
	void push_up(int p) {
		tree[p].lc = tree[ls].lc;
		tree[p].rc = tree[rs].rc;
		tree[p].sum = tree[ls].sum + tree[rs].sum;
		if (tree[ls].rc == tree[rs].lc) tree[p].sum--;
	}
	void push_down(int p) {
		if (tree[p].lazy) {
			down(ls, tree[p].lazy);
			down(rs, tree[p].lazy);
			tree[p].lazy = 0;
		}
	}
	
	void build(int p, int l, int r) {
		tree[p].l = l, tree[p].r = r;
		if (l == r) return;
		int mid = (l + r) >> 1;
		build(ls, l, mid);
		build(rs, mid + 1, r);
	}
	void modify(int p, int l, int r, int x) { // 區間賦值
		if (l <= tree[p].l && tree[p].r <= r) {
			down(p, x);
			return;
		}
		push_down(p);
		int mid = (tree[p].l + tree[p].r) >> 1;
		if (l <= mid) modify(ls, l, r, x);
		if (r > mid) modify(rs, l, r, x);
		push_up(p);
	}
	
	
	node query(int p, int l, int r) { // node 返回{左端點顏色，右端點顏色，顏色端數量}這個整體資訊
		if (l > tree[p].r || r < tree[p].l) return {-1, 0, 0};
		if (l <= tree[p].l && tree[p].r <= r) return {tree[p].lc, tree[p].rc, tree[p].sum};
		push_down(p);
		node x = query(ls, l, r), y = query(rs, l, r), res;
		if (x.lc == -1) return y;
		if (y.lc == -1) return x;
		res.lc = x.lc, res.rc = y.rc;
		res.sum = x.sum + y.sum - (x.rc == y.lc);
		return res;
	}
}tr;

考慮樹剖的過程，如果這次跳的路徑的底端和底端下面那個點顏色相同，則需要將顏色端數量減一。

int query_path(int u, int v) {
	int res = 0, t1 = -1, t2 = -1; // 維護t1,t2,代表當前u,v下方的顏色
	while (top[u] != top[v]) {
		if (dep[top[u]] < dep[top[v]]) swap(u, v), swap(t1, t2);
		node tmp = tr.query(1, id[top[u]], id[u]);
		// tmp.lc是top[u]的顏色
		// tmp.rc是u的顏色
		res += tmp.sum;
		if (t1 == tmp.rc) res--;
		t1 = tmp.lc; // 更新
		u = fat[top[u]];
	}
	if (dep[u] < dep[v]) swap(u, v), swap(t1, t2);
	node tmp = tr.query(1, id[v], id[u]);
	res += tmp.sum;
	if (tmp.rc == t1) res--;
	if (tmp.lc == t2) res--;
	return res;
}

CF916E Jamie and Tree

有一棵 \(n\) 個節點的有根樹，標號為 \(1\sim n\)，你需要維護以下三種操作

  1. 1 v：給定一個點 \(v\)，將整顆樹的根變為 \(v\)。

  2. 2 u v x：給定兩個點 \(u,v\) 和整數 \(x\)，將 \(\operatorname{lca}(u, v)\) 為根的子樹的所有點的點權都加上 \(x\)。

  3. 3 v：給定一個點 \(v\)，你需要回答以 \(v\) 所在的子樹的所有點的權值和。

這道題主要是加入了換根操作。設當前根節點為 \(root\)。

如果求 \(lca(u,v)\)，你只需要知道：\(lca(u,v)\) 等於 \(lca(u,root),lca(v,root),lca(u,v)\) 當中深度最大的那個點。具體證明略。

如果是對 \(u\) 子樹操作，需要分情況：

  • 如果 \(u=root\)，就是對整棵樹進行操作。

  • 如果 \(root\) 為 \(u\) 子樹，見下圖。

    image

    \(u\) 下方的靠近 \(root\) 的那個兒子叫做 \(t\)。可以這樣操作：先把整棵樹進行加，再把 t 子樹減回去。灰色部分就是被操作的點。

  • 如果 \(u\) 為 \(root\) 子樹，則直接對 \(u\) 進行操作即可。

#include <bits/stdc++.h>
using namespace std;
 
#define PII pair<int, int>
#define _for(i, a, b) for (int i = (a); i <= (b); i++)
#define _pfor(i, a, b) for (int i = (a); i >= (b); i--)
#define int long long
const int N = 3e5 + 5;
 
int n, m, a[N], sz[N], son[N], din[N], dep[N], idx, top[N], fat[N], id[N], cc, root;
vector<int> G[N];
 
struct fenwick {
	int c[N][2];
	void add(int x, int v) {
		for (int i = x; i <= N - 5; i += i & -i) {
			c[i][0] += v;
			c[i][1] += x * v;
		}
	}
	void modify(int l, int r, int v) {
		add(l, v);
		add(r + 1 , -v);
	}
	int sum(int op, int x) {
		int res = 0;
		for (int i = x; i; i -= i & -i) res += c[i][op];
		return res;
	}
	int query(int l, int r) {
		int t1 = sum(0, l - 1) * l - sum(1, l - 1);
		int t2 = sum(0, r) * (r + 1) - sum(1, r);
		return t2 - t1;
	}
}tr;
 
void dfs1(int u, int fa, int depth) {
	dep[u] = depth; sz[u] = 1; fat[u] = fa; din[u] = ++cc;
	for (auto v : G[u]) {
		if (v == fa) continue;
		dfs1(v, u, depth + 1);
		sz[u] += sz[v];
		if (sz[v] > sz[son[u]]) son[u] = v;
	}
}
 
void dfs2(int u, int fa, int tt) {
	top[u] = tt; id[u] = ++idx;
	if (!son[u]) return;
	dfs2(son[u], u, tt);
	for (auto v : G[u]) {
		if (v == fa || v == son[u]) continue;
		dfs2(v, u, v);
	}
}
 
int lca(int u, int v) {
	while (top[u] != top[v]) {
		if (dep[top[u]] < dep[top[v]]) swap(u, v);
		u = fat[top[u]];
	}
	if (dep[u] < dep[v]) return u;
	return v;
}
 
int get_son(int u, int v) { // 找v的兒子在u,v路徑上
	while (top[u] != top[v]) {
		if (fat[top[u]] == v) return top[u];
		u = fat[top[u]];
	}
	return son[v];
} 
 
int get_lca(int x, int y, int root) {
	PII t[3];
	t[0] = {dep[lca(x, y)], lca(x, y)};
	t[1] = {dep[lca(x, root)], lca(x, root)};
	t[2] = {dep[lca(y, root)], lca(y, root)};
	sort(t, t + 3);
	return t[2].second;
}
 
 
signed main() {
	cin >> n >> m;
	_for(i, 1, n) cin >> a[i];
	_for(i, 1, n - 1) {
		int u, v;
		cin >> u >> v;
		G[u].push_back(v);
		G[v].push_back(u);
	}
	dfs1(1, 0, 1);
	dfs2(1, 0, 1);
	_for(i, 1, n) tr.modify(din[i], din[i], a[i]);
	while (m--) {
		int op, u, v, x;
		cin >> op >> u;
		if (op == 1) root = u;
		else if (op == 2) {
			cin >> v >> x;
			int t = get_lca(u, v, root);
			if (t == root) tr.modify(1, n, x);
			else if (din[root] >= din[t] && din[root] <= din[t] + sz[t] - 1) {
				int tmp = get_son(root, t);
				tr.modify(1, n, x);
				tr.modify(din[tmp], din[tmp] + sz[tmp] - 1, -x);
			}
			else tr.modify(din[t], din[t] + sz[t] - 1, x);
		}
		else {
			int res = 0;
			if (u == root) res = tr.query(1, n);
			else if (din[root] >= din[u] && din[root] <= din[u] + sz[u] - 1) {
				int tmp = get_son(root, u);
				res += tr.query(1, n);
				res -= tr.query(din[tmp], din[tmp] + sz[tmp] - 1);
			}
			else res += tr.query(din[u], din[u] + sz[u] - 1);
			cout << res << endl;
		}
	}
}

相關文章