Wiki Cấu trúc dữ liệu Cây DSU on Tree (Small to Large)

DSU on Tree (Small to Large)

huunguyen huunguyen Updated Tháng sáu 2, 2026

DSU on Tree (còn gọi là Small to Large merging, hay biệt danh Sack) là kỹ thuật trả lời các truy vấn về cây con (subtree) cho mọi đỉnh trong tổng cộng O(NlogN) thay vì O(N2) như cách duyệt lại từng cây con một cách ngây thơ.

Bài toán mẫu kinh điển: cho cây N đỉnh, mỗi đỉnh có một màu; với mỗi đỉnh u, hãy đếm số màu phân biệt trong cây con gốc u. Nếu với mỗi đỉnh ta duyệt lại toàn bộ cây con của nó thì tốn O(N2). DSU on Tree khai thác cấu trúc cây để hạ xuống O(NlogN).

Điều kiện áp dụng: truy vấn offline (biết trước, không có cập nhật giá trị xen kẽ), và câu trả lời cho cây con của u chỉ phụ thuộc vào tập các đỉnh trong cây con đó (đếm phân biệt, mode, tần suất, số cặp thỏa điều kiện…), được duy trì qua một cấu trúc cộng/trừ dần (mảng đếm cnt[], BIT, set…).

Ý tưởng / Trực giác

Có hai biến thể, cùng một tinh thần "gộp tập nhỏ vào tập lớn".

Tại sao "small to large" lại nhanh?

Giả sử mỗi đỉnh giữ một tập (set) các phần tử của cây con nó. Khi xử lý đỉnh u, ta phải gộp tập của tất cả các con lại với nhau. Mẹo cốt lõi: luôn gộp tập nhỏ hơn vào tập lớn hơn (chứ không phải ngược lại).

Vì sao điều này cho O(NlogN)? Xét một phần tử bất kỳ. Mỗi lần nó bị di chuyển sang một tập khác, tập đích có kích thước ít nhất gấp đôi tập nguồn (vì ta chỉ chuyển khi tập nguồn nhỏ hơn tập đích). Một phần tử chỉ có thể "nhân đôi kích thước môi trường" tối đa log2N lần trước khi tập chứa nó đạt cỡ N. Vậy tổng số lần di chuyển trên toàn bộ là O(NlogN). Đây chính là lập luận kích thước nhân đôi quen thuộc của DSU.

Biến thể "heavy child" (DSU on Tree thuần)

Thay vì giữ set ở mỗi đỉnh, ta dùng một mảng đếm toàn cục cnt[] và một quy tắc thông minh dựa trên con nặng (heavy child) — con có cây con lớn nhất:

  1. Duyệt và giải quyết các con nhẹ (light) trước, mỗi con nhẹ tính xong thì xóa sạch đóng góp của nó khỏi cnt[].
  2. Duyệt con nặng cuối cùnggiữ lại đóng góp của nó trong cnt[].
  3. Thêm lại đóng góp của tất cả con nhẹ và của bản thân u vào cnt[], rồi ghi đáp án cho u.
  4. Nếu u là con nhẹ của cha nó, xóa toàn bộ cây con u khỏi cnt[] để cha tự xử lý.

Vì sao đúng và nhanh? Một đỉnh w bị thêm/xóa (add với sign = ±1) mỗi khi nó nằm trong cây con của một đỉnh mà cây con đó không phải con nặng trên đường đi từ w lên gốc. Số lần như vậy đúng bằng số cạnh nhẹ (light edge) trên đường đi từ w lên gốc, mà số cạnh nhẹ trên bất kỳ đường gốc–lá nào log2N (mỗi cạnh nhẹ làm kích thước cây con giảm ít nhất một nửa khi đi xuống). Vậy mỗi đỉnh bị đụng O(logN) lần tổng O(NlogN). Heavy child được "tái sử dụng miễn phí" nên không tốn thao tác xóa.

Ví dụ chạy tay

Xét cây sau, mỗi đỉnh ghi kèm màu của nó (ví dụ bài distcolor):

            1 (màu 2)
           /        \
        2 (3)       3 (2)
                   /     \
                4 (2)    5 (1)

Kích thước cây con: sz[1]=5, sz[2]=1, sz[3]=3, sz[4]=1, sz[5]=1. Con nặng: heavy[1]=3 (vì sz[3]=3 > sz[2]=1), heavy[3]=4 (4 và 5 cùng cỡ, chọn cái gặp trước).

Ta theo dõi mảng cnt[màu] (số lần mỗi màu xuất hiện trong tập đang xét) và distinct (số màu có cnt > 0).

Xử lý đỉnh 3 (gốc cây con {3,4,5} màu {2,2,1}):

  • Con nhẹ của 3 là đỉnh 5. Xử lý 5 (keep = false): thêm màu 1.
cnt: {1:1}        distinct = 1   →  ans[5] = 1
sau đó XÓA đỉnh 5:  cnt: {}        distinct = 0
  • Con nặng của 3 là đỉnh 4. Xử lý 4 (keep = true): thêm màu 2, giữ lại.
cnt: {2:1}        distinct = 1   →  ans[4] = 1   (giữ nguyên, không xóa)
  • Thêm lại con nhẹ (đỉnh 5, màu 1) và bản thân đỉnh 3 (màu 2):
+ màu 1:  cnt: {2:1, 1:1}   distinct = 2
+ màu 2:  cnt: {2:2, 1:1}   distinct = 2   →  ans[3] = 2
  • Đỉnh 3 là con nặng của 1 nên không xóa — để đỉnh 1 dùng lại.

Xử lý đỉnh 1 (gốc toàn bộ cây):

  • Con nhẹ của 1 là đỉnh 2. Xử lý 2 (keep = false): thêm màu 3.

Nhưng lúc này cnt vẫn còn {2:2, 1:1} từ con nặng (đỉnh 3). Để xử lý đỉnh 2 độc lập, ta gọi dfs(2, keep=false) trước khi giữ con nặng. Trật tự đúng là: light trước (và mỗi light tự dọn), rồi heavy giữ lại, rồi add lại light. Diễn lại cho đỉnh 1:

1) dfs các con nhẹ (đỉnh 2), keep=false:
      + màu 3:  cnt {3:1}      distinct=1  → ans[2]=1
      XÓA đỉnh 2: cnt {}       distinct=0
2) dfs con nặng (đỉnh 3), keep=true:  (như trên) cnt {2:2,1:1} distinct=2
3) add lại con nhẹ (đỉnh 2, màu 3) + bản thân đỉnh 1 (màu 2):
      + màu 3:  cnt {2:2,1:1,3:1}  distinct=3
      + màu 2:  cnt {2:3,1:1,3:1}  distinct=3  → ans[1]=3

Kết quả: ans = [3, 1, 2, 1, 1] — đúng như đáp án mẫu 3 1 2 1 1.

Cài đặt

Cách 1 — Small to Large với set (ngắn, dễ nhớ)

Phù hợp khi chỉ cần đếm số phần tử phân biệt. Đơn giản nhất nhưng kèm hệ số log của set.

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

int n;
vector<int> adj[200005];
int val[200005];          // màu của đỉnh
int ans[200005];
set<int>* S[200005];      // con trỏ tới set của mỗi đỉnh

void dfs(int u, int p) {
    S[u] = new set<int>();
    S[u]->insert(val[u]);
    for (int v : adj[u]) {
        if (v == p) continue;
        dfs(v, u);
        // Quan trọng: gộp set NHỎ vào set LỚN.
        // Nếu con lớn hơn, hoán đổi con trỏ để u trỏ vào set lớn.
        if (S[v]->size() > S[u]->size()) swap(S[u], S[v]);
        for (int x : *S[v]) S[u]->insert(x);
        delete S[v];                  // giải phóng để không tốn O(N^2) bộ nhớ
    }
    ans[u] = (int)S[u]->size();        // số màu phân biệt trong cây con u
}
Cách 2 — DSU on Tree (heavy child + mảng đếm toàn cục)

Nhanh hơn, linh hoạt hơn (mode, số cặp, tổng có điều kiện…). Đây là bản chuẩn cho bài distcolor.

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

const int MAXN = 200005;
int n;
vector<int> adj[MAXN];
int val[MAXN], comp[MAXN];     // comp = màu sau khi nén (compress)
int sz[MAXN], heavy[MAXN];
int cnt[MAXN];                 // cnt[c] = số lần màu c xuất hiện trong tập đang xét
int distinctNow = 0;           // số màu đang có cnt > 0
int ans[MAXN];

// Bước 1: tính kích thước cây con và con nặng
void dfs_sz(int u, int p) {
    sz[u] = 1; heavy[u] = -1;
    int best = 0;
    for (int v : adj[u]) {
        if (v == p) continue;
        dfs_sz(v, u);
        sz[u] += sz[v];
        if (sz[v] > best) { best = sz[v]; heavy[u] = v; }
    }
}

// Thêm (sign=+1) hoặc xóa (sign=-1) toàn bộ cây con u, BỎ QUA nhánh skip (con nặng đang giữ)
void update(int u, int p, int skip, int sign) {
    int c = comp[u];
    if (sign == 1) { if (cnt[c]++ == 0) distinctNow++; }
    else           { if (--cnt[c] == 0) distinctNow--; }
    for (int v : adj[u])
        if (v != p && v != skip) update(v, u, skip, sign);
}

// Bước 2: DSU on Tree
void dfs(int u, int p, bool keep) {
    // 1) Giải quyết các con NHẸ trước, mỗi con tự dọn sạch sau khi xong
    for (int v : adj[u])
        if (v != p && v != heavy[u]) dfs(v, u, false);
    // 2) Giải quyết con NẶNG, GIỮ LẠI đóng góp trong cnt[]
    if (heavy[u] != -1) dfs(heavy[u], u, true);
    // 3) Thêm lại các con nhẹ + bản thân u vào cnt[]
    //    (con nặng đã có sẵn nên truyền skip = heavy[u] để khỏi đếm lại)
    update(u, p, heavy[u], +1);
    ans[u] = distinctNow;          // ghi đáp án cho u
    // 4) Nếu u là con nhẹ, xóa toàn bộ cây con u để cha xử lý độc lập
    if (!keep) update(u, p, -1, -1);
}

int main() {
    scanf("%d", &n);
    vector<int> raw(n + 1);
    for (int i = 1; i <= n; i++) { scanf("%d", &raw[i]); val[i] = raw[i]; }
    // Nén màu về [0, n) vì màu có thể tới 1e9
    vector<int> srt(raw.begin() + 1, raw.end());
    sort(srt.begin(), srt.end());
    srt.erase(unique(srt.begin(), srt.end()), srt.end());
    for (int i = 1; i <= n; i++)
        comp[i] = lower_bound(srt.begin(), srt.end(), val[i]) - srt.begin();
    for (int i = 0; i < n - 1; i++) {
        int a, b; scanf("%d %d", &a, &b);
        adj[a].push_back(b); adj[b].push_back(a);
    }
    dfs_sz(1, 0);
    dfs(1, 0, true);
    for (int i = 1; i <= n; i++) printf("%d ", ans[i]);
    return 0;
}

Với cây sâu (chuỗi dài, N tới 2·105) đệ quy có thể tràn stack — nên tăng giới hạn stack hoặc chuyển sang DFS lặp khi cần.

Độ phức tạp

Phương pháp Thời gian Bộ nhớ
Small to Large với set O(Nlog2N) O(N)
Small to Large với unordered_set O(NlogN) trung bình O(N)
DSU on Tree (heavy child + mảng cnt) O(NlogN) O(N)

Vì sao thời gian như vậy?

  • Small to Large: mỗi phần tử bị di chuyển O(logN) lần (mỗi lần kích thước tập đích nhân đôi). Tổng O(NlogN) thao tác insert; nhân thêm O(logN) của set thành O(Nlog2N). Dùng unordered_set bỏ được một thừa số log.
  • DSU on Tree: mỗi đỉnh bị update(±1) đúng bằng số cạnh nhẹ trên đường lên gốc, log2N cạnh. Tổng O(NlogN) thao tác trên mảng cnt[] (mỗi thao tác O(1)). Con nặng được giữ lại nên không phát sinh thêm.

Bộ nhớ: Small to Large nếu không delete set con sau khi gộp sẽ phình tới O(N2) — nhớ giải phóng. DSU on Tree chỉ giữ một mảng cnt[] cỡ O(N) duy nhất.

⚠️ Lỗi thường gặp

  • Gộp lớn vào nhỏ (sai chiều). Quên if (S[v]->size() > S[u]->size()) swap(...) thì độ phức tạp tụt về O(N2) trong cây hình chuỗi — TLE dù logic vẫn ra đúng đáp án. Triệu chứng: AC trên test nhỏ, TLE test lớn.
  • Quên giải phóng/xóa tập con. Bản Small to Large mà không delete S[v] (hoặc DSU on tree quên update(-1) cho con nhẹ) làm bộ nhớ phình O(N2) hoặc đáp án bị "rò rỉ" giữa các cây con anh em — đỉnh sau nhận thừa số liệu của đỉnh trước.
  • Add trùng con nặng. Khi đã dfs(heavy, keep=true), lúc gọi update(u, +1) phải truyền skip = heavy[u] để không đếm lại cây con nặng. Quên skip làm tần suất bị nhân đôi, mode/đếm sai.
  • Không nén giá trị màu. Màu tới 109 mà dùng làm chỉ số mảng cnt[] trực tiếp gây tràn mảng / RE. Phải nén (coordinate compression) về [0,N) trước.
  • Tràn số khi cộng tổng. Nếu bài hỏi tổng giá trị cây con (mỗi giá trị tới 109, N tới 2·105), tổng có thể đạt ~2·1014 — phải dùng long long, dùng int sẽ tràn âm.
  • Tràn stack do đệ quy sâu. Cây hình chuỗi N=2·105 làm DFS đệ quy tràn stack (segfault). Cần nới stack hoặc viết DFS không đệ quy.
  • Quên thứ tự light-trước-heavy-sau. Nếu xử lý con nặng trước rồi mới đụng con nhẹ, cnt[] đã "bẩn" khi tính đáp án cho con nhẹ — kết quả con nhẹ sai hoàn toàn.

Biến thể / Mở rộng

  • Truy vấn cây con tổng quát: thay distinctNow bằng cấu trúc phù hợp — mảng tần suất cho mode, BIT cho "đếm phần tử lớn hơn / nhỏ hơn", tổng có điều kiện…
  • Đếm cặp thỏa điều kiện (ví dụ dating-like trên cây con, hoặc đếm cặp cùng giá trị): mỗi lần add một đỉnh thì cộng dồn số đỉnh đã có trong tập thỏa điều kiện với nó.
  • Small to Large trên cấu trúc khác set: gộp các map, các Treap, các đa thức (gộp đa thức nhỏ vào lớn) — cùng lập luận kích thước nhân đôi.
  • Khi có cập nhật/online: DSU on Tree không dùng được; chuyển sang Euler tour + BIT/Segment Tree (mỗi cây con là một đoạn liên tục trong dãy Euler) như bài subqry.

Bài tập luyện

  • Đếm đường đi qua đỉnh (cntpath)(Advanced) đếm số đường đi đi qua mỗi đỉnh; gộp small-to-large độ sâu các nhánh để đếm cặp endpoint, làm quen tư duy "gộp tập anh em".
  • Màu yêu thích (favcolor)(Advanced) bài gắn tag small_to_large: gộp các tập ràng buộc đồng màu, luyện kỹ thuật gộp nhỏ vào lớn trên đồ thị.
  • Số màu phân biệt trong cây con (distcolor)(Veteran) đúng bài mẫu trong trang này: với mỗi đỉnh đếm số màu phân biệt của cây con, áp dụng trực tiếp cả hai cách cài đặt.
  • Đếm cấp dưới xuất sắc (promote)(Master) với mỗi đỉnh đếm số con cháu có năng lực cao hơn nó: gộp small-to-large các tập giá trị (hoặc DSU on tree + BIT) trên cây con.
gnatmake 12.2.0 a68g 3.1.2 nasm 2.16.1 as_x64 2.46 awk 1.3.4 gcc 16.1.0 csc 6.12.0.200 g++ 16.1.0 g++-themis 16.1.0 g++17 16.1.0 g++20 16.1.0 g++23 16.1.0 clang++ 22.1.6 dmd 2.112.0 dart 3.12.1 gforth 0.7.3 gfortran 12.2.0 go 1.26.3 groovyc 5.0.6 javac 25.0.3 node 26.2.0 kotlinc 2.3.21 sbcl 2.2.9 lua 5.4.8 nim 2.2.10 fpc 3.2.2 fpc-themis 3.2.2 perl 5.36.0 php 8.5.6 pike 8.0 pypy3 7.3.23 python3 3.14.5 racket 8.7 ruby 4.0.5 rustc 1.96.0 csc 5.3.0 ctoj-scratch 0.0.1 sed 4.9 tclsh 8.6 bun 1.3.14 deno 2.8.1 v 0.5.1 zig 0.16.0