Quản lý nhóm bằng Disjoint Set Union

Lan tổ chức 6 bạn vào các nhóm dự án. Ban đầu mỗi bạn là một nhóm. Sau mỗi yêu cầu ghép hai bạn, cả hai nhóm của họ cùng làm việc. Lan cần trả lời nhanh hai bạn đã chung nhóm chưa và nhóm còn bao nhiêu người.

Mỗi người thuộc đúng một nhóm

Tập rời nhau (disjoint sets) là các nhóm không có phần tử chung. Nếu ghép nhóm của bạn 1 với nhóm của bạn 2, mọi người ở hai nhóm đều thành một nhóm. Ta không chỉ ghép riêng hai cá nhân.

Chữ A, B, C chỉ để nhận diện nhóm trong hình. Chúng không phải tên cố định của đại diện. Khi hai người trong cùng nhóm được ghép lại, các nhóm không đổi.

Cấu trúc tập rời nhau (Disjoint Set Union, viết tắt DSU) hỗ trợ hai thao tác chính. find(x) tìm đại diện nhóm của x. union(a, b) gộp hai nhóm chứa a và b. Hai phần tử cùng nhóm đúng khi hai kết quả find bằng nhau.

Mỗi phần tử lần theo cha để tìm đại diện

Ta lưu parent[x] là cha (parent) của x. Một gốc (root) là phần tử tự trỏ tới chính nó: parent[root] == root. Gốc là đại diện của nhóm. Các mũi tên cha tạo thành những cây riêng; bạn chỉ cần hiểu mũi tên từ một người dẫn về người đại diện.

Phần tử 4 trỏ tới 3, rồi 3 trỏ tới 1. Vì 1 tự trỏ, nhóm của 4 có đại diện 1. Phần tử 5 tự trỏ nên thuộc nhóm riêng. Ô số 0 không xuất hiện vì ví dụ đánh số từ 1.

Đoạn sau tìm đại diện mà chưa thay đổi mảng:

parent = [0, 1, 1, 1, 3, 5]
def find(x):
    while parent[x] != x:
        x = parent[x]
    return x
print(find(4), find(2), find(5))

1 1 5

Mỗi bước giữ ta trong cùng nhóm. Khi dừng tại gốc, ta đã tìm đúng đại diện. Nếu đường cha dài h, một lần tìm đi qua nhiều nhất h mũi tên. Không được để mũi tên tạo vòng giữa hai phần tử: vòng lặp sẽ không dừng.

Khi mới bắt đầu, mỗi phần tử tự làm gốc

Với n người đánh số từ 1, ta tạo n + 1 ô. Ô 0 không tính là người. Mỗi ô ban đầu lưu chính chỉ số của nó.

n = 5
parent = list(range(n + 1))
size = [1] * (n + 1)
print(parent[1:])
print(size[1:])

[1, 2, 3, 4, 5]
[1, 1, 1, 1, 1]

size[root] lưu số người của nhóm có đại diện root. Khi chưa gộp, mỗi nhóm có kích thước một. Nếu đánh số từ 0 đến n - 1, hãy dùng đúng n ô thay vì bỏ thêm ô 0.

Chỉ nối hai gốc khi gộp hai nhóm

Muốn ghép nhóm của a và b, trước tiên tìm hai gốc ra, rb. Nếu chúng bằng nhau, không có nhóm mới nào được tạo. Nếu khác nhau, cho một gốc trỏ tới gốc kia. Mọi phần tử dưới gốc cũ vẫn đi được tới gốc mới.

Hình dưới biểu diễn quan hệ cha trước và sau khi gộp. Mũi tên hướng về cha; ký hiệu gốc không có mũi tên đi ra.

Nếu gán parent[b] = a trực tiếp, b có thể không phải gốc. Khi đó ta có thể tách một phần nhóm cũ hoặc tạo vòng. Đoạn sau minh họa việc gộp đúng bằng gốc:

parent = list(range(6))
def find(x):
    while parent[x] != x:
        x = parent[x]
    return x
def union(a, b):
    a, b = find(a), find(b)
    if a == b:
        return False
    parent[b] = a
    return True
print(union(1, 2), union(3, 4), union(2, 3))
print(union(1, 4), find(4) == find(2))
print(parent[1:])

True True True
False True
[1, 1, 1, 3, 5]

Lần cuối trả False vì hai người đã chung nhóm. Kết quả đại diện phụ thuộc thứ tự nối. Điều cần so sánh là hai đại diện có bằng nhau, không phải đại diện có phải số nhỏ nhất hay không.

Gắn nhóm nhỏ dưới nhóm lớn để hạn chế chiều cao

Gộp theo kích thước (union by size) gắn gốc nhóm nhỏ vào gốc nhóm lớn. Nếu bằng nhau, chọn một trong hai gốc theo quy ước nhất quán. Sau đó cộng kích thước vào gốc còn sống.

parent = [0, 1, 1, 3, 3, 3, 6]
size = [1, 2, 1, 3, 1, 1, 1]
def find(x):
    while parent[x] != x:
        x = parent[x]
    return x
def union(a, b):
    a, b = find(a), find(b)
    if a == b:
        return False
    if size[a] < size[b]:
        a, b = b, a
    parent[b] = a
    size[a] += size[b]
    return True
union(2, 5)
r = find(2)
print(r, size[r], parent[1:])

3 5 [3, 1, 3, 3, 3, 6]

Nhóm gốc 1 có hai người, nhóm gốc 3 có ba người. Sau khi gộp, gốc 1 trỏ tới 3 và kích thước tại 3 là năm. size[1] còn giá trị cũ nhưng 1 không còn là gốc. Hỏi kích thước của x phải dùng size[find(x)].

Mỗi khi độ sâu của một phần tử tăng, nhóm chứa nó đã được gắn vào nhóm ít nhất lớn bằng nó. Kích thước nhóm vì thế ít nhất tăng gấp đôi. Với n = 100000, độ sâu không vượt khoảng 17 nếu chỉ dùng quy tắc này. Khởi tạo tốn O(n) bộ nhớ và thời gian; mỗi phép tìm hoặc gộp có cận O(log n).

Nén đường đi để lần tìm sau ngắn hơn

Nén đường đi (path compression) cho các phần tử vừa đi qua trỏ thẳng tới gốc. Quan hệ nhóm không đổi vì chúng đã có cùng gốc. Ta dùng hai vòng lặp, không cần đệ quy.

Vòng đầu tìm gốc. Vòng thứ hai đi lại đường cũ, nhớ cha cũ trước khi đổi mũi tên. Đoạn dưới in cả hai trạng thái để bạn quan sát:

parent = [0, 1, 1, 2, 3, 5]
def find(x):
    root = x
    while parent[root] != root:
        root = parent[root]
    while parent[x] != x:
        old = parent[x]
        parent[x] = root
        x = old
    return root
print(parent[1:])
print(find(4))
print(parent[1:])

[1, 1, 2, 3, 5]
1
[1, 1, 1, 1, 5]

Phải lưu old trước khi thay cha. Nếu chuyển x ngay sang cha mới, bạn có thể bỏ qua các nút trung gian. Một cách khác là nén từng phần bằng parent[x] = parent[parent[x]] khi đi lên.

parent = [0, 1, 1, 2, 3, 4]
def find(x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]
        x = parent[x]
    return x
print(find(5))
print(parent[1:])

1
[1, 1, 1, 3, 3]

Cách này không nhất thiết làm mọi nút trên đường cũ trỏ thẳng tới gốc ngay trong một lần. Nó vẫn giữ nguyên gốc và rút ngắn đường. Kết hợp nén với gộp theo kích thước, tổng thời gian của nhiều thao tác là O((n + q) α(n)). Ở đây α tăng cực chậm; đây là cận trung bình phân bổ trên cả chuỗi, không phải cam kết mỗi phép chỉ một bước.

Ghép các thao tác thành một cấu trúc dùng lại

Lớp sau quản lý đại diện và kích thước. Hàm union trả True đúng khi hai nhóm khác nhau đã gộp thành một.

class DSU:
    def __init__(self, n):
        self.parent = list(range(n + 1))
        self.size = [1] * (n + 1)
        self.groups = n

    def find(self, x):
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, a, b):
        a, b = self.find(a), self.find(b)
        if a == b:
            return False
        if self.size[a] < self.size[b]:
            a, b = b, a
        self.parent[b] = a
        self.size[a] += self.size[b]
        self.groups -= 1
        return True

d = DSU(6)
for a, b in [(1, 2), (3, 4), (2, 3), (1, 4)]:
    print(d.union(a, b), d.groups)
print(d.find(1) == d.find(4), d.size[d.find(4)])

True 5
True 4
True 3
False 3
True 4

Ban đầu có sáu nhóm. Mỗi lần gộp thành công thay hai nhóm bằng một nhóm, nên số nhóm giảm đúng một. Lần gộp lặp không thay đổi số nhóm hoặc kích thước. Mảng cha có thể thay đổi do nén ngay cả khi phép gộp không thành công; thành viên các nhóm vẫn không đổi.

Thông tin của nhóm phải chuyển cùng đại diện

Ngoài kích thước, một nhóm có thể lưu tổng điểm, số người đạt chuẩn hay mã nhỏ nhất. Chỉ gốc hiện tại có thông tin đáng tin. Khi gộp hai gốc khác nhau, kết hợp thông tin vào gốc được giữ lại.

parent = list(range(5))
size = [1] * 5
total = [0, 5, -2, 7, 4]
def find(x):
    while parent[x] != x:
        x = parent[x]
    return x
def union(a, b):
    a, b = find(a), find(b)
    if a == b:
        return
    if size[a] < size[b]:
        a, b = b, a
    parent[b] = a
    size[a] += size[b]
    total[a] += total[b]
union(1, 2)
union(3, 1)
print(total[find(2)], size[find(3)])

10 3

Điểm âm vẫn được cộng như bình thường. Nếu gộp lại hai người đã cùng nhóm, không được cộng tổng thêm lần nữa. Muốn chọn mã nhỏ nhất, lưu riêng minimum[root] và cập nhật bằng min; gốc DSU không nhất thiết là mã nhỏ nhất.

Đếm số cặp mới khi hai nhóm ghép lại

Giả sử hai nhóm rời nhau có a và b người. Mỗi người nhóm đầu có thể ghép với mỗi người nhóm sau. Vì vậy có a * b cặp mới chung nhóm. Các cặp cũ trong từng nhóm đã được tính từ trước.

sizes = [(1, 1), (2, 3), (4, 2)]
for a, b in sizes:
    print(a, b, a * b)

1 1 1
2 3 6
4 2 8

Khi hai người đã cùng nhóm, số cặp mới là 0. Trước khi gọi union, tìm hai gốc và đọc kích thước của chúng. Một phép nhân cho mỗi lần ghép thành công, nên chi phí thêm là O(1) mỗi thao tác.

Đọc yêu cầu ghép và trả lời truy vấn

Ví dụ sau nhận n, số thao tác q, rồi các dòng 1 a b để ghép hoặc 2 a b để hỏi cùng nhóm. Mỗi câu hỏi in YES hoặc NO.

5 5
1 1 2
2 1 3
1 2 3
2 1 3
2 4 4

import sys
data = list(map(int, sys.stdin.read().split()))
n, q = data[0], data[1]
parent = list(range(n + 1))
size = [1] * (n + 1)
def find(x):
    while parent[x] != x:
        parent[x] = parent[parent[x]]
        x = parent[x]
    return x
out = []
for i in range(q):
    kind, a, b = data[2 + 3*i:5 + 3*i]
    a, b = find(a), find(b)
    if kind == 1 and a != b:
        if size[a] < size[b]:
            a, b = b, a
        parent[b] = a
        size[a] += size[b]
    elif kind == 2:
        out.append('YES' if a == b else 'NO')
print('\n'.join(out))

NO
YES
YES

Khi hỏi cùng một người, hai đại diện hiển nhiên bằng nhau. Năm thao tác chỉ tạo ba dòng vì thao tác ghép không in. Với 100000 người và 100000 thao tác, hai mảng có khoảng 200002 ô. Thời gian xử lý gần tuyến tính theo số thao tác, ngoài chi phí đọc và in.

DSU lưu quan hệ nhóm, không lưu cách nối

DSU phù hợp khi các nhóm chỉ được ghép thêm. Bạn có thể hỏi cùng nhóm, số nhóm hoặc thông tin đã tổng hợp ở gốc. Cấu trúc này không lưu danh sách những yêu cầu ghép ban đầu. Nó cũng không tự tách nhóm khi một yêu cầu bị hủy.

Ví dụ, cả hai chuỗi ghép (1, 2), (2, 3) và (1, 3), (1, 2) đều tạo nhóm gồm ba người. Từ kết luận ba người cùng nhóm, ta không thể khôi phục chắc chắn chuỗi yêu cầu. Nếu cần lịch sử, lưu riêng các yêu cầu; nếu cần xóa liên kết, cần kỹ thuật khác.

histories = [[(1, 2), (2, 3)], [(1, 3), (1, 2)]]
for history in histories:
    groups = [{1}, {2}, {3}]
    for a, b in history:
        left = next(g for g in groups if a in g)
        right = next(g for g in groups if b in g)
        if left is not right:
            left.update(right)
            groups.remove(right)
    print(sorted(groups[0]))

[1, 2, 3]
[1, 2, 3]

Ví dụ này dùng tập hợp để làm rõ ý nghĩa, không phải cách cài DSU cho dữ liệu lớn. Các ứng dụng dùng cạnh, chu trình và cây khung được học khi đã có nền đồ thị. Ở đây, hãy thành thạo thao tác gộp nhóm trước.

Lỗi thường gặp

Tóm tắt

• find tìm đại diện bằng cách đi theo cha; hai phần tử cùng nhóm khi đại diện bằng nhau.
• union nối hai gốc khác nhau. Gộp theo kích thước giúp cây thấp; nén đường đi giúp các lần tìm sau ngắn hơn.
• Chỉ gốc hiện tại giữ kích thước và thông tin nhóm chính xác. Gộp lặp không đổi số nhóm.
• DSU phù hợp với việc ghép dần, nhưng không tự lưu lịch sử hay hỗ trợ tách nhóm.

Kiểm tra nhanh

Đọc thêm

• Python: định nghĩa hàm.
• Python: cấu trúc dữ liệu danh sách.