Bỏ qua

DSU trên cây

Giới thiệu

Thuật toán "heuristic" là gì?

Thuật toán heuristic là các kỹ thuật tối ưu hóa dựa trên kinh nghiệm và trực giác của con người.

Ví dụ điển hình nhất là kỹ thuật hợp nhất theo heuristic trong cấu trúc dữ liệu hợp nhất tập hợp (DSU), mã như sau:

1
2
3
4
5
6
void merge(int x, int y) {
  int xx = find(x), yy = find(y);
  if (size[xx] < size[yy]) swap(xx, yy);
  fa[yy] = xx;
  size[xx] += size[yy];
}

Ở đây, với hai tập hợp có kích thước khác nhau, ta sẽ hợp tập nhỏ vào tập lớn, thay vì ngược lại.

Tại sao lại như vậy? Kích thước tập hợp có thể coi là chiều cao của cây (trong trường hợp thông thường), và việc hợp cây thấp vào cây cao sẽ giúp việc tìm cha nhanh hơn.

Việc cho cây thấp làm con của cây cao là một tối ưu gọi là hợp nhất theo heuristic.

Nội dung thuật toán

Thuật toán "dsu on tree" (hợp nhất theo heuristic trên cây) là một kỹ thuật mạnh để giải các bài toán truy vấn offline trên cây, vừa dễ hiểu vừa dễ cài đặt, tốc độ ngang hoặc vượt nhiều thuật toán khác.

Xét bài toán sau: Đếm số màu trên cây.

Ví dụ dẫn nhập

Cho một cây \(n\) đỉnh gốc tại \(1\), mỗi đỉnh \(u\) có màu \(c_u\). Với mỗi đỉnh \(u\), hỏi trong cây con gốc \(u\) có bao nhiêu màu khác nhau.

\(n\le 2\times 10^5\).

dsu-on-tree-1.png

Với dạng bài này, thường phải dùng các cấu trúc dữ liệu phức tạp (như cây phân đoạn lồng nhau), nhưng nếu cho phép truy vấn offline, liệu có cách đơn giản hơn?

Quy trình

Vì cho phép truy vấn offline, ta có thể tiền xử lý để trả lời truy vấn trong \(O(1)\).

Nếu duyệt brute-force, độ phức tạp là \(O(n^2)\): với mỗi đỉnh, duyệt toàn bộ cây con, mỗi lần \(O(n)\), tổng \(O(n^2)\).

Nhận thấy đáp án của mỗi đỉnh phụ thuộc vào cây con của nó và chính nó, ta có thể tận dụng tính chất này.

Trước tiên, tiền xử lý kích thước cây con và xác định "con nặng" (giống như trong phân tích chuỗi nặng nhẹ - HLD), con nặng là con có cây con lớn nhất, quá trình này \(O(n)\).

Gọi \(cnt_i\) là số lần màu \(i\) xuất hiện, \(ans_u\) là đáp án của đỉnh \(u\).

Duyệt một đỉnh \(u\) theo các bước sau:

  1. Duyệt các con nhẹ (không phải con nặng), tính đáp án nhưng không giữ lại ảnh hưởng lên mảng \(cnt\);
  2. Duyệt con nặng, giữ lại ảnh hưởng lên mảng \(cnt\);
  3. Duyệt lại cây con của các con nhẹ, cộng thêm ảnh hưởng để tính đáp án cho \(u\).

dsu-on-tree-2.png

Hình trên minh họa quy trình.

Như vậy, mỗi đỉnh chỉ duyệt cây con nặng một lần, cây con nhẹ hai lần, rất tối ưu.

Sau khi thực hiện, ta có đáp án cho mọi cây con.

Tại sao không gộp bước 1 và 3? Vì mảng \(cnt\) không thể dùng lại, nếu không sẽ tốn quá nhiều bộ nhớ, cần đảm bảo dùng \(O(n)\) không gian.

Rõ ràng, nếu một đỉnh \(u\) bị duyệt \(x\) lần, thì con nặng bị duyệt \(x\) lần, con nhẹ (nếu có) bị duyệt \(2x\) lần.

Lưu ý: sau khi duyệt xong con nhẹ, phải reset mảng \(cnt\).

Chứng minh

Giống như phân tích chuỗi nặng nhẹ, định nghĩa cạnh nặng và nhẹ (cạnh nối tới con nặng là cạnh nặng, còn lại là cạnh nhẹ). Xem hình dưới, với cây \(n\) đỉnh:

Số cạnh nhẹ từ gốc tới bất kỳ đỉnh không quá \(\log n\). Gọi số cạnh nhẹ là \(x\), kích thước cây con là \(y\), rõ ràng cây con nối qua cạnh nhẹ nhỏ hơn một nửa cha (nếu lớn hơn thì là con nặng), nên \(y<n/2^x\), tức \(n>2^x\), nên \(x<\log n\).

Nếu một đỉnh là con nặng của cha, thì cây con của nó lớn nhất trong các anh em, nên trên đường từ gốc tới nó, các cha nối qua cạnh nặng sẽ không duyệt lại nó, nên số lần bị duyệt là số cạnh nhẹ trên đường + 1 (vì chính nó cũng bị duyệt), tổng số lần bị duyệt là \(\log n+1\), tổng thời gian \(O(n(\log n+1))=O(n\log n)\), trả lời truy vấn \(O(m)\).

dsu-on-tree-3.png

Các cạnh tô đậm là cạnh nặng, con nối qua cạnh nặng là con nặng

Tối ưu

Như đã nói, dsu on tree tận dụng khái niệm con nặng/nhẹ để tăng tốc. Ta cũng có thể dùng luôn thứ tự dfs từ phân tích chuỗi nặng nhẹ để chuyển đệ quy thành lặp, tối ưu thêm hằng số.

Thứ tự dfs có tính chất: cây con của một đỉnh là một đoạn liên tiếp trên thứ tự dfs. Do đó, có thể duyệt ngược thứ tự dfs, đảm bảo khi duyệt một đỉnh, cây con của nó đã được xử lý.

Thứ tự dfs từ phân tích chuỗi nặng nhẹ còn có tính chất: một chuỗi nặng là một đoạn liên tiếp trên thứ tự dfs. Khi duyệt ngược, nếu đỉnh là đầu chuỗi nặng, thì đỉnh tiếp theo không phải cha nó, nên phải reset ảnh hưởng; nếu không phải đầu chuỗi nặng, thì đỉnh trước là con nặng hoặc đã reset, nên có thể kế thừa ảnh hưởng. Dựa vào đó, dùng thứ tự dfs để cộng nhanh ảnh hưởng của các con nhẹ, ghi lại đáp án.

Quy trình này gọi là dsu on tree không đệ quy/lặp (hay dsu on tree theo thứ tự dfs). So với bản đệ quy, nó giảm chi phí gọi hàm và bộ nhớ, đặc biệt hiệu quả với cây nhiều chuỗi, tiết kiệm stack.

Cài đặt

Mã tham khảo
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
#include <cstdio>
#include <vector>
using namespace std;

constexpr int N = 2e5 + 5;

int n, m;

// g[u]: 存储与 u 相邻的结点
vector<int> g[N];

// sz: 子树大小
// big: 重儿子
// col: 结点颜色
// L[u]: 结点 u 的 DFS 序
// R[u]: 结点 u 子树中结点的 DFS 序的最大值
// Node[i]: DFS 序为 i 的结点
// totdfn: 节点计数器,也是当前遍历过节点的 DFS 序最大值
// ans: 存答案
// cnt[i]: 颜色为 i 的结点个数
// totColor: 目前出现过的颜色个数
int sz[N], big[N], col[N], L[N], R[N], Node[N], totdfn;
int ans[N], cnt[N], totColor;

void add(int u) {
  if (cnt[col[u]] == 0) ++totColor;
  cnt[col[u]]++;
}

void del(int u) {
  cnt[col[u]]--;
  if (cnt[col[u]] == 0) --totColor;
}

int getAns() { return totColor; }

void dfs0(int u, int p) {
  L[u] = ++totdfn;
  Node[totdfn] = u;
  sz[u] = 1;
  for (int v : g[u])
    if (v != p) {
      dfs0(v, u);
      sz[u] += sz[v];
      if (!big[u] || sz[big[u]] < sz[v]) big[u] = v;
    }
  R[u] = totdfn;
}

void dfs1(int u, int p, bool keep) {
  // 计算轻儿子的答案
  for (int v : g[u])
    if (v != p && v != big[u]) {
      dfs1(v, u, false);
    }
  // 计算重儿子答案并保留计算过程中的数据(用于继承)
  if (big[u]) {
    dfs1(big[u], u, true);
  }
  for (int v : g[u])
    if (v != p && v != big[u]) {
      // 子树结点的 DFS 序构成一段连续区间,可以直接遍历
      for (int i = L[v]; i <= R[v]; i++) {
        add(Node[i]);
      }
    }
  add(u);
  ans[u] = getAns();
  if (!keep) {
    for (int i = L[u]; i <= R[u]; i++) {
      del(Node[i]);
    }
  }
}

int main() {
  scanf("%d", &n);
  for (int i = 1; i < n; i++) {
    int u, v;
    scanf("%d%d", &u, &v);
    g[u].push_back(v);
    g[v].push_back(u);
  }
  for (int i = 1; i <= n; i++) scanf("%d", &col[i]);
  dfs0(1, 0);
  dfs1(1, 0, false);
  scanf("%d", &m);
  for (int i = 1; i <= m; i++) {
    int q;
    scanf("%d", &q);
    printf("%d\n", ans[q]);
  }
  return 0;
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
#include <cstdio>
#include <vector>
using namespace std;

constexpr int N = 2e5 + 5;

int n, m;

// g[u]: 存储与 u 相邻的节点
vector<int> g[N];

// sz: 子树大小
// son[u]: 节点u的重儿子的编号
// col: 节点颜色
// dfn[u]: 节点 u 的 DFS 序
// bottom[u]: 节点 u 子树中节点的 DFS 序的最大值
// totdfn: 节点计数器,也是当前遍历过节点的 DFS 序最大值
// fa[u]: 节点u的父亲编号
// top[u]:节点u所在重链的顶端节点编号
// rnk[i]: DFS 序为 i 的节点
// ans[u]: 存答案
// cnt[i]: 颜色为 i 的节点个数
// totColor: 目前出现过的颜色个数
int sz[N], son[N], col[N], dfn[N], fa[N], bottom[N], totdfn;
int top[N], rnk[N];
int ans[N], cnt[N], totColor;

void add(int u) {
  if (cnt[col[u]] == 0) ++totColor;
  cnt[col[u]]++;
}

void del(int u) {
  cnt[col[u]]--;
  if (cnt[col[u]] == 0) --totColor;
}

int getAns() { return totColor; }

void dfs0(int u, int p) {
  sz[u] = 1;
  fa[u] = p;
  for (int v : g[u])
    if (v != p) {
      dfs0(v, u);
      sz[u] += sz[v];
      if (!son[u] || sz[son[u]] < sz[v]) son[u] = v;
    }
}

void dfs1(int u, int t) {
  rnk[bottom[u] = dfn[u] = ++totdfn] = u;
  top[u] = t;
  if (!son[u]) return;
  dfs1(son[u], t);
  for (int v : g[u])
    if (v != fa[u] && v != son[u]) dfs1(v, v);
  bottom[u] = totdfn;
}

void dsu_on_tree() {
  for (int i = totdfn; i >= 1; i--) {
    int u = rnk[i];
    for (int v : g[u])
      if (v != fa[u] && v != son[u]) {
        for (int i = dfn[v]; i <= bottom[v]; i++) {
          add(rnk[i]);
        }
      }
    add(u);
    ans[u] = getAns();
    if (u == top[u]) {
      for (int i = dfn[u]; i <= bottom[u]; i++) {
        del(rnk[i]);
      }
    }
  }
}

int main() {
  scanf("%d", &n);
  for (int i = 1; i < n; i++) {
    int u, v;
    scanf("%d%d", &u, &v);
    g[u].push_back(v);
    g[v].push_back(u);
  }
  for (int i = 1; i <= n; i++) scanf("%d", &col[i]);
  dfs0(1, 1);
  dfs1(1, 1);
  dsu_on_tree();
  scanf("%d", &m);
  for (int i = 1; i <= m; i++) {
    int q;
    scanf("%d", &q);
    printf("%d\n", ans[q]);
  }
  return 0;
}

Ứng dụng

  1. Một số bài chính xác yêu cầu dùng dsu on tree

    Ví dụ CF741D: Cho một cây, mỗi đỉnh mang ký tự từ 'a' đến 'v', mỗi truy vấn hỏi trong cây con có đường đi nào mà các ký tự trên đường đi (sắp xếp lại) tạo thành chuỗi đối xứng.

    Vì chuỗi đối xứng nên mỗi ký tự xuất hiện hai lần coi như không xuất hiện, tức là đường đi thỏa mãn tối đa một ký tự xuất hiện lẻ.

    Cách thường là dfs từng đỉnh, mỗi lần duyệt liệt kê mọi ký tự, kiểm tra số ký tự lẻ lớn hơn 1 thì loại, lấy đường đi dài nhất, độ phức tạp \(O(n^2\log n)\), dùng dsu on tree tối ưu xuống \(O(n\log^2n)\). Chi tiết xem phần mở rộng bên dưới.

  2. Một số bài có thể dùng dsu để "ăn gian"

    Có thể dùng dsu để giải một số bài cây phân đoạn (không có truy vấn cập nhật), dsu có độ phức tạp tốt hơn Mo's trên cây \(O(n\sqrt{m})\).

Bài tập

CF600E Lomsat gelral

Dịch đề: Mỗi đỉnh trên cây có màu, một màu chiếm cây con khi không có màu nào xuất hiện nhiều hơn nó trong cây con. Tính tổng các màu chiếm mỗi cây con.

UOJ284 Gà chơi game vui vẻ

CF1709E XOR Tree

Tài liệu tham khảo/mở rộng

Blog tác giả CF741D về dsu on tree

Blog giải thích chi tiết