Biến cây thành đoạn và tìm tổ tiên chung gần nhất

Lan quản lý một sơ đồ gồm bảy trạm liên lạc. Mỗi trạm ngoài trạm gốc có một trạm cha. Lan muốn cộng số vật tư của cả nhánh và tìm nơi hai tuyến gặp nhau. Nếu duyệt lại toàn bộ cây cho từng câu hỏi, công việc có thể lặp rất nhiều.

Thứ tự đi vào đặt mỗi cây con trên một đoạn liên tiếp

Cây bên dưới có gốc 1. Các cạnh hướng xuống con, còn thứ tự xét các con là từ trái sang phải.

Một cây con (subtree) của u gồm u và mọi đỉnh đi xuống từ u. Cây con của 2 gồm 2, 4, 5. Trong phép duyệt sâu, ta vào 2, thăm hết hai nhánh của nó rồi mới trở lại 1. Vì vậy ba đỉnh này đứng liền nhau trong thứ tự vào.

Thứ tự Euler lúc vào (Euler entry order) chỉ ghi mỗi đỉnh một lần, ngay khi gặp đỉnh. Có cách gọi Euler tour ghi đỉnh nhiều lần khi quay lui. Hai cách có mục đích khác nhau; ở đây ta dùng cách một lần.

Chỉ số của 2 là 1. Mốc ngay sau nhánh là 4. Ta ký hiệu tin[2] = 1 và tout[2] = 4. Khoảng nửa mở [tin[2], tout[2]) chứa chỉ số 1, 2, 3. Nó không chứa chỉ số 4.

Trong đoạn dưới đây, ngăn xếp chứa một sự kiện vào hoặc ra. Mỗi sự kiện ra được đặt trước các sự kiện vào của con. Vì ngăn xếp lấy phần tử cuối trước, sự kiện ra chỉ được xử lý sau mọi con.

g = [[], [2, 3], [1, 4, 5], [1, 6], [2], [2], [3, 7], [6]]
tin = [0] * 8
tout = [0] * 8
order = []
stack = [(1, 0, False)]
while stack:
    u, parent, leaving = stack.pop()
    if leaving:
        tout[u] = len(order)
        continue
    tin[u] = len(order)
    order.append(u)
    stack.append((u, parent, True))
    for v in reversed(g[u]):
        if v != parent:
            stack.append((v, u, False))
print(order)
print(tin[2], tout[2], order[tin[2]:tout[2]])

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

Mỗi cạnh cây được xét hai lần qua danh sách kề. Mỗi đỉnh có một sự kiện vào và một sự kiện ra. Do đó thời gian và bộ nhớ là O(n) với n đỉnh. Với 100000 đỉnh, số lượt xét cạnh là 199998. Ngăn xếp tường minh cũng tránh giới hạn độ sâu đệ quy của Python trên cây hình dây.

Tại một lá, mốc ra lớn hơn mốc vào đúng một. Với cây chỉ gồm gốc, khoảng của gốc là [0, 1). Không tăng bộ đếm khi ra: tout ở đây là độ dài hiện tại của order.

order = [1]
tin = 0
tout = len(order)
print(tin, tout, order[tin:tout])

0 1 [1]

Tổng tiền tố trả lời tổng vật tư của một nhánh

Gán vật tư theo đỉnh: 1:5, 2:3, 3:8, 4:-2, 5:4, 6:1, 7:6. Số âm có thể biểu thị lượng đã dùng. Xếp các giá trị theo thứ tự vào để nhận một dãy.

order = [1, 2, 4, 5, 3, 6, 7]
value = [0, 5, 3, 8, -2, 4, 1, 6]
flat = [value[u] for u in order]
print(flat)

[5, 3, -2, 4, 8, 1, 6]

Ta tính tổng tiền tố (prefix sum) pref[i] của i phần tử đầu. Tổng một đoạn [l, r) bằng pref[r] - pref[l]. Đoạn cây con của 2 là [1, 4), nên không cộng phần tử ở vị trí 4.

flat = [5, 3, -2, 4, 8, 1, 6]
pref = [0]
for x in flat:
    pref.append(pref[-1] + x)
tin, tout = 1, 4
print(pref[tout] - pref[tin])

5

Mỗi giá trị trong đoạn xuất hiện một lần trong pref[r] và không xuất hiện trong pref[l]. Giá trị trước đoạn xuất hiện ở cả hai tổng nên bị trừ đi. Lập dãy và tiền tố mất O(n); mỗi câu hỏi mất O(1). Với 100000 trạm và 10000 câu hỏi, không cần duyệt lại 100000 trạm cho từng câu hỏi.

Nếu giá trị thay đổi, tiền tố cũ không còn đúng. Khi đó giữ nguyên thứ tự vào và cập nhật vị trí tin[u] trong Fenwick. Bài 12 giải thích cấu trúc này. Mỗi lần cập nhật điểm và tính tổng một đoạn mất O(log n).

Ví dụ sau dùng lại cây bảy trạm và đổi giá trị của trạm 4 từ -2 thành 7. Vì đây là phép đặt giá trị, ta cộng hiệu 7 - (-2) vào Fenwick. Mốc của cây con 2 vẫn là [1, 4).

flat = [5, 3, -2, 4, 8, 1, 6]
n = len(flat)
bit = [0] * (n + 1)
def add(pos, delta):
    i = pos + 1
    while i <= n:
        bit[i] += delta
        i += i & -i
def prefix(end):
    total = 0
    while end > 0:
        total += bit[end]
        end -= end & -end
    return total
for i, value in enumerate(flat):
    add(i, value)
print(prefix(4) - prefix(1))
position = 2
new_value = 7
add(position, new_value - flat[position])
flat[position] = new_value
print(prefix(4) - prefix(1))

5
14

Lần hỏi thứ hai phản ánh giá trị mới của trạm 4. Chỉ vị trí của trạm đó thay đổi; các mốc của cây vẫn giữ nguyên. Nếu cộng trực tiếp 7 thay cho hiệu, kết quả sẽ sai vì giá trị cũ chưa được loại bỏ.

Khoảng vào và ra cho biết quan hệ tổ tiên

Đỉnh a là tổ tiên (ancestor) của b nếu b thuộc cây con của a. Theo quy ước này, một đỉnh cũng là tổ tiên của chính nó. Vì cây con là đoạn liên tiếp, điều kiện tương đương là tin[a] <= tin[b] < tout[a].

tin = [0, 0, 1, 4, 2, 3, 5, 6]
tout = [0, 7, 4, 7, 3, 4, 7, 7]
def ancestor(a, b):
    return tin[a] <= tin[b] < tout[a]
print(ancestor(2, 5))
print(ancestor(2, 3))
print(ancestor(2, 2))

True
False
True

Vị trí vào của 5 nằm trong khoảng của 2. Vị trí vào của 3 đúng bằng mốc ra của 2, nên bị loại bởi dấu <. Điều kiện dùng hai phép so sánh, vì thế mỗi câu hỏi tốn O(1) sau tiền xử lý O(n).

Nếu đề hỏi tổ tiên nghiêm ngặt, phải thêm điều kiện a != b. Đổi gốc có thể đổi quan hệ tổ tiên. Các mốc trên chỉ đúng với gốc đã chọn lúc duyệt.

Tổ tiên chung gần nhất là điểm gặp sâu nhất của hai đường lên

Tổ tiên chung gần nhất (lowest common ancestor, LCA) của u và v là tổ tiên chung có độ sâu lớn nhất. Độ sâu của gốc là 0; mỗi cạnh đi xuống tăng độ sâu một. Trong cây của Lan, LCA của 4 và 5 là 2. LCA của 4 và 7 là 1.

Trường hợp u là tổ tiên của v, đáp án là chính u. Khi hai đỉnh trùng nhau, đáp án cũng là đỉnh đó. Có thể tìm bằng cách đi từng cha, nhưng đường dây dài 100000 đỉnh khiến mỗi câu hỏi tốn tới O(n).

parent = [0, 1, 1, 1, 2, 2, 3, 6]
depth = [0, 0, 1, 1, 2, 2, 2, 3]
u, v = 4, 7
while depth[u] > depth[v]:
    u = parent[u]
while depth[v] > depth[u]:
    v = parent[v]
while u != v:
    u, v = parent[u], parent[v]
print(u)

1

Hai đỉnh sau khi cân bằng độ sâu cùng đi lên từng cạnh. Trước khi gặp nhau, chúng chưa thể ở trên LCA. Lần gặp đầu tiên chính là tổ tiên chung sâu nhất. Cách này đơn giản để kiểm tra dữ liệu nhỏ, nhưng cần tăng tốc khi có nhiều câu hỏi.

Bảng nhảy nhị phân ghép các bước có độ dài lũy thừa hai

Nhảy nhị phân (binary lifting) lưu up[k][u], tổ tiên nhận được sau 2^k bước từ u. Ta quy ước cha của gốc là chính gốc. Nhảy quá gốc theo bảng sẽ ở gốc; nếu đề yêu cầu báo không tồn tại thì cần kiểm tra số bước với độ sâu trước.

Hàng đầu là up[0][u] = parent[u]. Hai lần nhảy 2^(k-1) ghép thành một lần nhảy 2^k. Vì vậy up[k][u] = up[k-1][up[k-1][u]].

parent = [0, 1, 1, 1, 2, 2, 3, 6]
up = [parent]
for k in range(1, 3):
    prev = up[k - 1]
    up.append([prev[prev[u]] for u in range(8)])
print(up[0][7], up[1][7], up[2][7])

6 3 1

Từ 7, một bước tới 6, hai bước tới 3, bốn bước tới gốc theo quy ước. Để đi lên ba bước, biểu diễn 3 = 2 + 1 và dùng hàng 1 rồi hàng 0. Mỗi bit bật tương ứng với một lần nhảy.

up = [[0, 1, 1, 1, 2, 2, 3, 6]]
for k in range(1, 3):
    prev = up[-1]
    up.append([prev[prev[u]] for u in range(8)])
u, steps = 7, 3
for k in range(3):
    if steps & (1 << k):
        u = up[k][u]
print(u)

1

Với n đỉnh, n.bit_length() hàng bao phủ mọi khoảng cách tối đa n - 1. Mỗi hàng tính n giá trị. Thời gian và bộ nhớ tiền xử lý là O(n log n); một lần nhảy theo số bước mất O(log n).

Tìm LCA bằng cách cân bằng độ sâu rồi nhảy cùng nhau

Trước hết đưa đỉnh sâu hơn lên cùng độ sâu với đỉnh kia. Nếu hai đỉnh trùng nhau, đó là LCA. Nếu chưa trùng, xét các mức từ lớn đến nhỏ. Khi hai tổ tiên ở cùng mức nhảy khác nhau, đưa cả hai lên mức đó.

Hai đỉnh đang có cùng độ sâu. Nếu hai tổ tiên sau một bước nhảy trùng, đưa cả hai lên sẽ tới hoặc vượt LCA. Nếu chúng khác nhau, cả hai vẫn nằm dưới LCA. Sau khi xét mọi mức, hai đỉnh khác nhau có cùng cha là LCA.

Chương trình sau đọc cây không trọng số, gốc 1, và trả lời nhiều câu hỏi. Cây có thể là dây dài; phép duyệt dùng ngăn xếp thay cho đệ quy.

7 4
1 2
1 3
2 4
2 5
3 6
6 7
4 5
4 7
3 7
6 6

import sys
data = list(map(int, sys.stdin.read().split()))
it = iter(data)
n, q = next(it), next(it)
g = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    a, b = next(it), next(it)
    g[a].append(b)
    g[b].append(a)
parent = [0] * (n + 1)
depth = [0] * (n + 1)
parent[1] = 1
stack = [1]
while stack:
    u = stack.pop()
    for v in g[u]:
        if v != parent[u]:
            parent[v] = u
            depth[v] = depth[u] + 1
            stack.append(v)
log = n.bit_length()
up = [parent]
for k in range(1, log):
    prev = up[k - 1]
    up.append([prev[prev[u]] for u in range(n + 1)])
def lca(u, v):
    if depth[u] < depth[v]:
        u, v = v, u
    diff = depth[u] - depth[v]
    for k in range(log):
        if diff & (1 << k):
            u = up[k][u]
    if u == v:
        return u
    for k in range(log - 1, -1, -1):
        if up[k][u] != up[k][v]:
            u, v = up[k][u], up[k][v]
    return up[0][u]
answers = []
for _ in range(q):
    answers.append(str(lca(next(it), next(it))))
print('\n'.join(answers))

2
1
3
6

Với 100000 đỉnh, log bằng 17. Mỗi lần hỏi xét nhiều nhất 17 mức khi cân bằng và 17 mức khi tìm điểm gặp. Tổng thời gian là O((n + q) log n) và bộ nhớ O(n log n). Python có chi phí bộ nhớ cho các số và danh sách; cần so giới hạn bộ nhớ của đề.

Khoảng cách và đỉnh trên đường đi dùng chung bảng LCA

Gọi w là LCA của u, v. Đường duy nhất từ u tới v đi lên tới w, rồi đi xuống. Vì thế số cạnh là depth[u] + depth[v] - 2 * depth[w].

depth = [0, 0, 1, 1, 2, 2, 2, 3]
u, v, w = 4, 7, 1
distance = depth[u] + depth[v] - 2 * depth[w]
print(distance)

5

Với cạnh có trọng số, thay độ sâu bằng tổng trọng số từ gốc. Công thức vẫn đúng vì hai đoạn từ gốc tới w bị tính hai lần. Số cạnh và tổng trọng số là hai đại lượng khác nhau; không thay thế lẫn nhau.

Để tìm đỉnh cách u đúng k cạnh trên đường tới v, đặt a = depth[u] - depth[w] và b = depth[v] - depth[w]. Nếu k <= a, nhảy lên k bước từ u. Nếu k > a, nhảy lên a + b - k bước từ v. Đề cần đảm bảo 0 <= k <= a + b.

depth_u, depth_v, depth_w = 2, 3, 0
k = 4
a = depth_u - depth_w
b = depth_v - depth_w
if k <= a:
    print('u', k)
else:
    print('v', a + b - k)

v 1

Ví dụ này chọn phía v và đi ngược một cạnh. Tìm LCA và một lần nhảy đều mất O(log n), nên mỗi câu hỏi vẫn mất O(log n).

Tổng giá trị trên đường đi cần giữ lại giá trị tại điểm gặp

Giả sử mỗi trạm có một giá trị, và bạn muốn cộng tất cả trạm trên tuyến từ u đến v. Gọi path[u] là tổng các giá trị từ gốc tới u, tính cả hai đầu. Khác với độ sâu, path[u] có thể giảm nếu giá trị âm.

Xét cây gồm các cạnh 1–2, 1–3, 2–4, 2–5. Năm giá trị theo số hiệu là 5, -2, 3, 4, 1. Giá trị âm biểu thị vật tư đã dùng. Chương trình sau tính tổng từ gốc cho mỗi đỉnh.

parent = [0, 1, 1, 1, 2, 2]
order = [1, 2, 4, 5, 3]
value = [0, 5, -2, 3, 4, 1]
path = [0] * 6
for u in order:
    path[u] = value[u]
    if u != 1:
        path[u] += path[parent[u]]
print(path[1:])

[5, 3, 8, 7, 4]

Cha xuất hiện trước con trong thứ tự vào. Do đó, khi tính path[u], tổng của cha đã sẵn sàng. Phần tính tổng này đi qua mỗi đỉnh một lần, tốn O(n) thời gian và O(n) bộ nhớ.

Gọi w là LCA của u và v. Tổng path[u] + path[v] đếm hai lần mọi đỉnh từ gốc tới w. Ta trừ 2 * path[w] để bỏ phần chung, rồi cộng lại value[w] đúng một lần. Công thức là path[u] + path[v] - 2 * path[w] + value[w].

value = [0, 5, -2, 3, 4, 1]
path = [0, 5, 3, 8, 7, 4]
u, v, w = 3, 5, 1
answer = path[u] + path[v] - 2 * path[w] + value[w]
print(answer)

7

Đường cần tính gồm 3, 1, 2, 5. Đỉnh 1 là điểm gặp và vẫn thuộc đường, nên phải được cộng trở lại. Nếu quên bước đó, đáp án thiếu giá trị ở w. Khi hai đầu trùng nhau, công thức còn đúng bằng giá trị của đỉnh ấy.

value = [0, 5, -2, 3, 4, 1]
path = [0, 5, 3, 8, 7, 4]
u = v = w = 2
print(path[u] + path[v] - 2 * path[w] + value[w])

-2

Với q tuyến và cây tĩnh, lập bảng nhảy mất O(n log n). Mỗi tuyến cần một lần tìm LCA và vài phép cộng trừ, mất O(log n). Không nên cộng lại từng đỉnh trên mỗi tuyến dài. Chẳng hạn, 8000 tuyến dài 50000 cạnh có thể yêu cầu khoảng 400 triệu lượt đi từng cạnh.

Nếu giá trị nằm trên cạnh thay vì trên đỉnh, tổng từ gốc tới w bị trừ hai lần là đủ. Lúc đó không cộng value[w]. Hãy đọc kỹ đề để biết giá trị thuộc đỉnh hay cạnh, và hai đầu có được tính hay không.

Chọn cách xử lý theo loại câu hỏi và thay đổi

Không phải câu hỏi nào trên cây cũng cần cùng một bảng. Nếu chỉ có một câu hỏi LCA và cây nhỏ, đi từng cha có thể dễ kiểm tra hơn. Với nhiều câu hỏi, chi phí tiền xử lý bảng nhảy thường được chia đều cho các lần hỏi.

Các kết quả trên giả sử các cạnh và gốc không đổi. Nếu thêm hoặc xóa cạnh, khoảng cây con và bảng tổ tiên cũ có thể không còn đúng. Nếu chỉ đổi giá trị đỉnh, cấu trúc cây và bảng LCA vẫn đúng. Tuy nhiên, các tổng tiền tố và tổng từ gốc đã lưu có thể cần cập nhật.

Khi viết chương trình, hãy thử ba dạng cây. Cây một đỉnh kiểm tra gốc và câu hỏi hai đầu trùng nhau. Cây sao kiểm tra hai lá có chung gốc. Cây dây kiểm tra chênh lệch độ sâu lớn và tránh đệ quy quá sâu.

Lỗi thường gặp

Tóm tắt

• Duyệt sâu theo thứ tự vào đưa cây con vào [tin[u], tout[u]).
• Tổng tiền tố trả lời tổng cây con tĩnh trong O(1) sau O(n) tiền xử lý.
• Bảng up kết hợp các bước lũy thừa hai để tìm LCA trong O(log n).
• Khoảng cách và vị trí trên đường đi suy ra từ LCA và độ sâu.

Kiểm tra nhanh

Đọc thêm

• Lowest Common Ancestor — Binary Lifting
• Tài liệu Python: kiểu số nguyên