Hai dòng code biến union-find từ O(n) thành gần O(1)

Hình dung một đám đông chia thành nhiều câu lạc bộ, và cách duy nhất để biết ai cùng câu lạc bộ với ai là: mỗi người chỉ tay vào người đã rủ mình vào, và bạn cứ đi theo chuỗi "ai rủ ai" đó tới khi gặp người sáng lập — đại diện của câu lạc bộ. Hai người cùng người sáng lập thì cùng câu lạc bộ. Nghe đơn giản, nhưng nếu chuỗi giới thiệu dài dằng dặc — A rủ B, B rủ C, C rủ D... — thì mỗi lần hỏi "câu lạc bộ nào" bạn phải lần ngược cả trăm nghìn người. Đó chính là union-find, và cái chuỗi dài kia là cái bẫy khiến nó chậm.

Union-find (disjoint set union, DSU) theo dõi các nhóm: find(x) trả về "đại diện" nhóm của x, union(a, b) gộp hai nhóm. Nó là xương sống của thuật toán Kruskal (cây khung nhỏ nhất), của việc tìm thành phần liên thông, của rất nhiều bài toán "phần tử này có cùng nhóm với phần tử kia không". Cách cài trông đơn giản đến mức nguy hiểm: một mảng parent[], mỗi phần tử trỏ tới cha, find đi lên tới gốc. Nhưng nếu bạn dừng ở đó, cấu trúc có thể chậm thảm họa. Tôi đo bản ngây thơ và bản tối ưu trong container gcc:13, và con số cho thấy hai dòng code tối ưu là ranh giới giữa O(n) và gần O(1).

Đi lên gốc — và cái bẫy chuỗi dài

Cài cơ bản: parent[i] là cha của i, gốc trỏ tới chính nó. find(x) đi lên chuỗi cha tới khi gặp gốc. union(a, b) tìm hai gốc rồi gắn gốc này dưới gốc kia. Nghe xong xuôi.

Vấn đề nằm ở hình dạng cây. Nếu union cứ gắn tùy tiện, cây có thể suy biến thành một chuỗi dài — phần tử 0 dưới 1, 1 dưới 2, …, N phần tử thành một sợi thẳng cao N. Lúc đó find phải đi hết chuỗi = O(n) mỗi lần, đúng như cây tìm kiếm suy biến ở phần 7, đúng cái chuỗi giới thiệu dài dằng dặc ở đầu bài. Cấu trúc "O(1) trong tưởng tượng" trở thành O(n) trong thực tế.

Hai tối ưu kinh điển sửa điều này:

• Nén đường (path compression): mỗi lần find(x), sau khi tìm ra gốc, kéo mọi node trên đường trỏ thẳng tới gốc — nên lần sau find chỉ đi một bước. Cây tự làm phẳng theo thời gian. (Giống như: sau khi lần ra người sáng lập một lần, mọi người trên chuỗi ghi nhớ luôn người sáng lập, lần sau hỏi là biết ngay.)
• Hợp theo hạng (union by rank/size): khi gộp hai cây, luôn gắn cây thấp dưới cây cao (hoặc cây nhỏ dưới cây lớn) — để chiều cao không tăng vô ích. Cây luôn nông.

Kết hợp cả hai, mỗi thao tác gần như O(1) — chính xác là O(α(n)), với α là hàm Ackermann ngược, một hàm tăng chậm đến mức α(n) < 5 với mọi n thực tế (kể cả n bằng số nguyên tử trong vũ trụ). Nói cách khác: gần như hằng số.

Đo: từ 28.368 ns/find xuống 3,8 ns/thao tác

Tôi dựng DSU với 1 triệu phần tử, làm 2 triệu union + 2 triệu find, so các phiên bản:

N=1 triệu phần tử, 2 triệu union + 2 triệu find (g++ -O2):

OPT (nén đường + union by rank):
   3,8 ns/thao tác     (~O(α(n)), gần O(1))

CHỈ nén đường (union tùy tiện):
   14,1 ns/thao tác    (vẫn rất nhanh)

NAIVE (không nén, union tạo chuỗi dài):
   chuỗi 200k phần tử, 20k find: 28.367,8 ns/find   (O(n)! find đi hết chuỗi)

Nhìn bản NAIVE: khi để cây suy biến thành chuỗi, mỗi find mất 28.367,8 nano-giây — vì nó phải đi qua hàng trăm nghìn node. Đây là O(n) đầy đủ, và tôi phải giảm quy mô (chuỗi 200k, chỉ 20k find) để nó không chạy mãi. So với bản OPT (nén đường + union by rank) chỉ 3,8 ns/thao tác — nhanh hơn ~7.500 lần mỗi thao tác, và đó là còn tính OPT trên khối lượng lớn hơn nhiều. Bản chỉ nén đường đã đạt 14,1 ns/thao tác — cho thấy nén đường một mình gần như đủ để tránh thảm họa, nhưng thêm union by rank vẫn nhanh hơn ~4 lần nữa.

Điều tôi thấy đáng nhớ nhất: cả ba đều là "union-find", cùng ý tưởng, cùng mảng parent[]. Khác biệt duy nhất là hai tối ưu vài dòng — và chúng biến một cấu trúc O(n) thành một cấu trúc nhanh hơn cả O(log n).

(Môi trường đo: container gcc:13 trên host ARM AArch64, g++ -O2, N=1 triệu, mỗi số lấy MIN sau warmup. Con số đổi theo máy; cái cần nhìn là vực O(n) khi cây suy biến so với gần-hằng-số khi có hai tối ưu.)

Hai niềm tin đối nghịch, cùng sai

Tôi vào đo với niềm tin: "union-find đơn giản, cứ đi lên gốc cây là xong". Con số bác thẳng: bản ngây thơ, khi union tạo chuỗi, cho find O(n) — 28.368 ns/find, chậm như duyệt một danh sách dài. Hai tối ưu (nén đường + union by rank) không phải trang trí — chúng là điều bắt buộc để union-find nhanh như tiếng tăm của nó.

Nhưng có một niềm tin ngược, tinh vi hơn: "vậy chắc nó là O(log n), như cây cân bằng". Cũng sai — với cả hai tối ưu, union-find đạt gần O(1) amortize, O(α(n)), nhanh hơn O(log n). Nó là một trong số hiếm cấu trúc "gần như hằng số thật sự" nhờ một phân tích amortize sâu. Đánh giá thấp nó là dùng cấu trúc khác chậm hơn cho đúng những bài mà DSU là lựa chọn tối ưu.

Về đến code hằng ngày

Điều đầu tiên: luôn cài cả hai tối ưu — chúng chỉ vài dòng và là bắt buộc. Nén đường: trong find, sau khi tìm gốc, cho các node trỏ thẳng tới gốc (hoặc dùng "path halving" p[x] = p[p[x]] — đơn giản hơn, gần như hiệu quả bằng). Union by rank/size: giữ một mảng rank hoặc size, gắn cây nhỏ dưới cây lớn. Đừng bỏ qua "cho gọn code" — cái giá là O(n) ẩn. Đây là mẫu quen thuộc suốt sê-ri: một tối ưu nhỏ ngăn cấu trúc suy biến về trường hợp tệ nhất.

Điều thứ hai: biết union-find là công cụ "gần O(1)" cho bài toán nhóm/liên thông. Khi bạn cần "gộp nhóm và hỏi cùng nhóm không" — thành phần liên thông động, Kruskal, phát hiện chu trình, gom cụm — union-find gần như luôn là lựa chọn tốt nhất, nhanh hơn cả cây cân bằng hay hash. Nhận ra bài toán có hình dạng đó và với tới DSU là một kỹ năng đáng có.

Chỗ này nếu bạn nhìn ra hệ thống thật sẽ thấy DSU chạy ở những nơi bất ngờ. Trình biên dịch dùng union-find trong suy luận kiểu (type inference kiểu Hindley-Milner): mỗi lần hợp nhất hai biến kiểu là một union, và thuật toán này là lý do compiler của ML/Haskell/Rust suy được kiểu nhanh đến vậy. Các thuật toán xử lý ảnh gán nhãn vùng liên thông (connected-component labeling) dùng DSU để gộp các mảng pixel. Và trong database phân tán hay hệ thống mạng, "các node này có thuộc cùng một partition/cluster không" cũng là bài toán disjoint-set. Điểm chung: bất cứ đâu bạn chỉ gộp chứ không tách các nhóm và cần hỏi thành viên liên tục, union-find với hai tối ưu là câu trả lời gần-như-miễn-phí.

Muốn tự thấy thì cài union-find bằng một mảng parent[] với find đi lên gốc và union gắn tùy tiện. Rồi cố tình tạo chuỗi: union(0,1), union(1,2), union(2,3), ... cho vài trăm nghìn phần tử, và bấm giờ một find ở cuối chuỗi — nó chậm khủng khiếp, vì đi hết chuỗi. Giờ thêm nén đường (trong find, cho các node trỏ thẳng tới gốc) và đo lại chính find đó — lần thứ hai gần như tức thời, vì đường đã bị nén. Thêm union by rank và làm lại thí nghiệm với hàng triệu thao tác: bạn sẽ thấy mỗi thao tác chỉ vài nano-giây, gần như hằng số dù N lớn cỡ nào. Vài phút đó cho bạn thấy điều mà "đi lên gốc là xong" giấu đi: union-find chỉ đạt tốc độ huyền thoại của nó nhờ hai tối ưu giữ cây phẳng — thiếu chúng, nó suy biến thành một danh sách liên kết chậm chạp mang tên cấu trúc dữ liệu nhanh nhất.