Tính nhanh tổng đoạn và cập nhật dữ liệu

Lan theo dõi số sách mượn trong sáu ngày: 3, 1, 4, 2, 5 và 2. Cô cần trả lời hàng nghìn câu hỏi về tổng của những đoạn ngày khác nhau. Nếu cộng lại từng ngày cho từng câu hỏi, cùng một số có thể bị cộng rất nhiều lần.

Tổng tiền tố giúp trả lời tổng đoạn bằng phép trừ

Một tiền tố (prefix) gồm các phần tử đầu tiên của dãy. Quy ước p[0] = 0 và p[i] là tổng của i phần tử đầu tiên. Khi đó tổng đoạn nửa kín [l, r) bằng p[r] - p[l].

Đoạn nửa kín chứa vị trí l nhưng không chứa r. Nó có r - l phần tử và có thể rỗng khi hai đầu bằng nhau.

Với l = 1, r = 4, hiệu là 10 - 3 = 7. Các phần tử được tính là 1, 4 và 2. Số 3 đứng trước đoạn nên xuất hiện trong cả hai tiền tố và bị trừ đi.

a = [3, 1, 4, 2, 5, 2]
p = [0]
for x in a:
    p.append(p[-1] + x)
for l, r in [(1, 4), (0, 6), (3, 3)]:
    print(p[r] - p[l])

7
17
0

Vòng lặp tạo p trong O(n) thời gian và dùng O(n) bộ nhớ. Mỗi câu hỏi sau đó cần một phép trừ, tức O(1). Với 100 000 phần tử và 100 000 câu hỏi, cách cộng lại có thể cần 10 tỉ lượt cộng. Tạo tiền tố rồi hỏi chỉ cần cỡ 200 000 bước chính.

Nếu đề đánh số từ 1 và hỏi đoạn đóng [L, R], kết quả là p[R] - p[L - 1]. Hai quy ước đều đúng, miễn bạn chuyển nhất quán.

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

n, q = map(int, input().split())
a = list(map(int, input().split()))
p = [0]
for x in a:
    p.append(p[-1] + x)
for _ in range(q):
    left, right = map(int, input().split())
    print(p[right] - p[left - 1])

8
12
2

Số âm không làm công thức mất hiệu lực. Phép trừ loại đúng phần đứng trước đoạn, bất kể tổng có tăng hay giảm.

a = [5, -8, 3, -1]
p = [0]
for x in a:
    p.append(p[-1] + x)
print(p)
print(p[4] - p[1])

[0, 5, -3, 0, -1]
-6

Đoạn một phần tử và đoạn phủ cả dãy là hai phép thử ngắn nhưng hiệu quả. Với đoạn đóng từ 1 đến n, công thức cho p[n] - p[0]. Với đoạn chỉ chứa vị trí i, hiệu p[i] - p[i-1] phải trả lại đúng giá trị ô đó.

a = [4, -3, 6, -2]
p = [0]
for value in a:
    p.append(p[-1] + value)
for left, right in [(1, 4), (1, 1), (4, 4)]:
    print(p[right] - p[left - 1])

5
4
-2

Kết quả đầu là tổng cả bốn số. Hai kết quả sau lần lượt lấy riêng số đầu và số cuối. Các kiểm tra này giúp phát hiện lỗi lệch chỉ số, kể cả khi tổng ở giữa dãy tình cờ đúng.

Mảng hiệu ghi lại nơi một lần cộng bắt đầu và kết thúc

Giả sử mỗi học sinh trong các vị trí từ l tới r nhận thêm v điểm. Nếu có nhiều lượt cập nhật và chỉ cần kết quả cuối, cộng từng vị trí sẽ lặp lại công việc. Mảng hiệu (difference array) ghi thay đổi ở hai biên.

Với đoạn đóng đánh số từ 0, cộng v vào d[l] và trừ v tại d[r + 1]. Duyệt tổng tích lũy của d sẽ khôi phục số cần cộng vào từng vị trí. Cấp n + 1 ô để cả trường hợp r = n - 1 đều có chỗ ghi biên kết thúc.

a = [2, 2, 2, 2, 2]
d = [0] * (len(a) + 1)
for left, right, value in [(1, 3, 3), (0, 4, -1)]:
    d[left] += value
    d[right + 1] -= value
running = 0
for i in range(len(a)):
    running += d[i]
    a[i] += running
print(*a)

1 4 4 4 1

Mỗi cập nhật chỉ sửa hai ô, nên m cập nhật và một lượt khôi phục tốn O(n + m). Với 100 000 ô và 100 000 cập nhật, cách này cần khoảng 300 000 thao tác biên và duyệt. Mảng hiệu không tự trả lời giá trị mới tức thời sau mỗi cập nhật.

Với mảng hiệu, hãy thử một cập nhật kết thúc tại phần tử cuối. Ô chặn sau cuối không xuất hiện trong kết quả, nhưng vẫn phải có chỗ lưu. Cũng cần thử hai cập nhật triệt tiêu nhau và cập nhật mang giá trị âm.

Đếm đoạn có tổng bằng k bằng cách đếm tiền tố đã gặp

Xét đoạn kết thúc ngay trước vị trí r. Đoạn [l, r) có tổng k khi p[l] = p[r] - k. Ta lưu số lần xuất hiện của mỗi tổng tiền tố trước đó trong một từ điển.

Trước khi đọc phần tử đầu, tiền tố 0 đã xuất hiện một lần. Ở mỗi bước, đếm tiền tố cần tìm trước, rồi mới ghi tiền tố hiện tại. Thứ tự này loại đoạn rỗng khi k = 0.

a = [1, -1, 3, 0, 2]
k = 3
counts = {0: 1}
prefix = answer = 0
for x in a:
    prefix += x
    answer += counts.get(prefix - k, 0)
    counts[prefix] = counts.get(prefix, 0) + 1
print(answer)

4

Mỗi tiền tố trước đó tương ứng với một vị trí bắt đầu riêng. Vì vậy cộng tần suất sẽ đếm đủ các đoạn, kể cả khi tổng tiền tố lặp lại. Thời gian kỳ vọng là O(n) với thao tác từ điển trung bình O(1); bộ nhớ là O(n). Với n = 100 000, ta duyệt 100 000 phần tử thay vì xét gần 5 tỉ đoạn.

a = [1, -1, 3, 0, 2]
k = 3
answer = 0
for l in range(len(a)):
    total = 0
    for r in range(l, len(a)):
        total += a[r]
        answer += total == k
print(answer)

4

Cách vét cạn hữu ích để kiểm tra dữ liệu nhỏ. Nó cần O(n²) lượt mở rộng đoạn nên không phù hợp với giới hạn lớn.

Với bài đếm đoạn, hãy thử dãy chỉ gồm số 0. Hai số 0 tạo ba đoạn khác rỗng: hai đoạn một phần tử và một đoạn hai phần tử. Nếu đếm được năm, chương trình có thể đã tính thêm hai đoạn rỗng.

a = [0, 0]
freq = {0: 1}
s = fast = 0
for x in a:
    s += x
    fast += freq.get(s, 0)
    freq[s] = freq.get(s, 0) + 1
slow = 0
for left in range(len(a)):
    for right in range(left + 1, len(a) + 1):
        slow += sum(a[left:right]) == 0
print(fast, slow)

3 3

Hai cách phải cho cùng số đoạn. Cách vét cạn chỉ là công cụ đối chiếu trên dãy ngắn; nó không thay thế cách đếm tiền tố với đầu vào lớn.

Xét dãy [2, -2, 2, -2] và đích 0. Các tiền tố lần lượt là 0, 2, 0, 2, 0. Hai tiền tố bằng nhau xác định một đoạn khác rỗng ở giữa chúng.

Tiền tố 0 xuất hiện ba lần, tạo ba cách chọn hai vị trí. Tiền tố 2 xuất hiện hai lần, tạo thêm một cách. Tổng cộng có bốn đoạn, bao gồm cả đoạn phủ toàn bộ dãy.

a = [2, -2, 2, -2]
for left in range(len(a)):
    total = 0
    for right in range(left, len(a)):
        total += a[right]
        if total == 0:
            print(left + 1, right + 1)

1 2
1 4
2 3
3 4

Chương trình liệt kê vị trí đầu và cuối theo chỉ số từ 1. Với đích khác 0, hai tiền tố không cần bằng nhau: tiền tố trước cần bằng tiền tố hiện tại trừ đi đích. Đó là lý do tra freq[s-k] thay vì chỉ tra freq[s].

Từ điển lưu số lần gặp chứ không chỉ lưu có hay không. Nếu cùng một tiền tố xuất hiện ba lần trước vị trí hiện tại, cả ba vị trí đều có thể mở đầu một đoạn hợp lệ. Ghi mỗi giá trị một lần sẽ bỏ sót các đoạn đó.

Tổng tiền tố hai chiều loại phần bị đếm thừa

Với bảng a gồm n hàng và m cột, p[i][j] là tổng hình chữ nhật gồm i hàng đầu và j cột đầu. Hàng 0 và cột 0 của p đều bằng 0. Công thức cộng ô mới là p[i][j] = a[i-1][j-1] + p[i-1][j] + p[i][j-1] - p[i-1][j-1].

Vùng góc trên trái được cộng trong cả hai tiền tố nên phải trừ đi một lần. Tổng hình chữ nhật nửa kín [r1, r2) × [c1, c2) tính tương tự bằng bốn giá trị.

a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
n, m = len(a), len(a[0])
p = [[0] * (m + 1) for _ in range(n + 1)]
for i in range(1, n + 1):
    for j in range(1, m + 1):
        p[i][j] = a[i-1][j-1] + p[i-1][j] + p[i][j-1] - p[i-1][j-1]
r1, c1, r2, c2 = 1, 1, 3, 3
print(p[r2][c2] - p[r1][c2] - p[r2][c1] + p[r1][c1])

28

Tạo bảng tốn O(nm) thời gian và bộ nhớ. Mỗi câu hỏi chỉ cần bốn ô, tức O(1). Với bảng 1 000 × 1 000, cần khoảng một triệu ô tiền tố, nên phải kiểm tra cả giới hạn bộ nhớ.

Một bảng hai chiều cần quy ước rõ hàng và cột trước khi tính. Xét khu vườn có hai hàng và ba cột, mỗi số là lượng nước đo được ở một ô. Bạn cần tổng của hai cột cuối trong cả hai hàng.

garden = [[2, 1, 3], [4, -2, 5]]
rows, cols = 2, 3
p = [[0] * (cols + 1) for _ in range(rows + 1)]
for i in range(1, rows + 1):
    for j in range(1, cols + 1):
        p[i][j] = (garden[i-1][j-1] + p[i-1][j]
                   + p[i][j-1] - p[i-1][j-1])
for row in p:
    print(*row)
print(p[2][3] - p[0][3] - p[2][1] + p[0][1])

0 0 0 0
0 2 3 6
0 6 5 13
7

Phần cần lấy gồm các giá trị 1, 3, -2 và 5. Lấy toàn bộ hai hàng rồi trừ cột đầu cho cùng kết quả. Trong trường hợp tổng quát, khi trừ vùng trên và vùng trái, góc chung bị trừ hai lần và cần cộng lại.

Đừng khởi tạo bằng [[0] * (m + 1)] * (n + 1). Phép nhân danh sách ngoài lặp lại tham chiếu đến cùng một hàng. Sửa một ô sẽ thay đổi ô cùng cột ở các hàng khác.

shared = [[0] * 3] * 2
separate = [[0] * 3 for _ in range(2)]
shared[0][1] = 8
separate[0][1] = 8
print(shared)
print(separate)

[[0, 8, 0], [0, 8, 0]]
[[0, 8, 0], [0, 0, 0]]

Ở danh sách thứ nhất, cả hai hàng cùng thay đổi. Ở danh sách thứ hai, mỗi lần lặp tạo một hàng mới nên chỉ hàng đầu thay đổi. Dùng cách thứ hai khi tạo bảng tiền tố.

Cây Fenwick hỗ trợ cập nhật điểm xen kẽ truy vấn

Một mảng tiền tố thông thường phải sửa nhiều ô khi một giá trị thay đổi. Cây Fenwick (Fenwick tree, hay binary indexed tree) lưu tổng các khối để hỗ trợ cả cập nhật điểm và truy vấn tiền tố. Cấu trúc dùng chỉ số từ 1 tới n.

Đặt lowbit(i) = i & -i. Ô tree[i] chứa tổng của lowbit(i) phần tử kết thúc tại i. Chẳng hạn tree[6] chứa vị trí 5 và 6 vì lowbit(6) = 2. Ô tree[4] chứa vị trí 1 tới 4.

for i in range(1, 9):
    size = i & -i
    print(i, size, i - size + 1, i)

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

Để tính tổng từ 1 tới i, cộng khối tại i rồi chuyển sang i - lowbit(i). Các khối này kề nhau và không trùng, cho tới khi chỉ số bằng 0. Để cộng vào một vị trí, sửa mọi khối chứa nó bằng cách liên tục chuyển sang i + lowbit(i).

Ví dụ dưới đây gom mảng cây và hai thao tác vào một lớp (class). Đây là cú pháp mới so với khóa Cơ bản, nhưng bạn chỉ cần đọc được bốn điểm sau:

• Fenwick(6) tạo một đối tượng (object) mới. Python gọi __init__ để chuẩn bị dữ liệu ban đầu cho đối tượng đó.
• self là chính đối tượng đang được dùng. self.tree là danh sách gắn với đối tượng, nên vẫn còn giữ giá trị giữa các lần gọi.
• def add(self, i, delta) là một hàm thuộc lớp, còn gọi là phương thức. Khi viết bit.add(3, 2), Python tự truyền bit vào vị trí self.
• Nếu không dùng lớp, bạn vẫn có thể viết tree = [0] * (n + 1) cùng hai hàm add(tree, i, delta) và prefix(tree, i). Lớp chỉ giúp gọi gọn hơn.

class Fenwick:
    def __init__(self, n):
        self.n = n
        self.tree = [0] * (n + 1)

    def add(self, i, delta):
        while i <= self.n:
            self.tree[i] += delta
            i += i & -i

    def prefix(self, i):
        result = 0
        while i > 0:
            result += self.tree[i]
            i -= i & -i
        return result

bit = Fenwick(6)
for i, value in enumerate([3, 1, 4, 2, 5, 2], 1):
    bit.add(i, value)
print(bit.prefix(5) - bit.prefix(1))
bit.add(3, -2)
print(bit.prefix(5) - bit.prefix(1))

12
10

Mỗi bước truy vấn bỏ bit thấp nhất đang bật; mỗi bước cập nhật chuyển tới khối lớn hơn tiếp theo. Do đó mỗi thao tác tốn O(log n), còn mảng cây dùng O(n) bộ nhớ. Với 100 000 phần tử, mỗi thao tác cần cỡ tối đa 17 bước. Cách khởi tạo bằng n lần add tốn O(n log n).

Không gọi add(0, delta): 0 & -0 bằng 0, khiến vòng lặp không tiến. Nếu đề yêu cầu gán giá trị mới, hãy lưu giá trị cũ và cập nhật hiệu mới - cũ.

a = [3, 1, 4]
old = a[1]
new = 7
delta = new - old
a[1] = new
print(delta, sum(a))

6 14

Giả sử dãy có tám vị trí. Khi hỏi tổng đến vị trí 7, khối đầu chứa riêng vị trí 7. Chuyển sang 6 cho khối 5–6, rồi sang 4 cho khối 1–4.

Ba khối kề nhau, không bỏ sót và không đếm trùng vị trí nào. Do mỗi bước bỏ bit thấp nhất đang bật, truy vấn đi đến 0 sau không quá số bit của n.

Khi cộng vào vị trí 3, ta sửa các khối kết thúc ở 3, 4 và 8. Mỗi khối ấy chứa vị trí 3. Với n = 8, chỉ số tiếp theo là 16 nên ta dừng.

n = 8
i = 3
updates = []
while i <= n:
    updates.append(i)
    i += i & -i
i = 7
queries = []
while i > 0:
    queries.append(i)
    i -= i & -i
print(*updates)
print(*queries)

3 4 8
7 6 4

Một lần cộng làm thay đổi đúng các khối chứa vị trí đó. Các khối không chứa nó vẫn giữ tổng cũ. Bởi vậy, sau mọi cập nhật, mỗi ô vẫn giữ đúng ý nghĩa ban đầu.

Với n = 100 000, số bit cần biểu diễn chỉ số là 17. Một truy vấn hoặc cập nhật đi qua tối đa khoảng 17 ô, thay vì sửa hay cộng lại hàng chục nghìn ô. Hệ số thực tế còn phụ thuộc cách cài đặt và máy chạy.

Chọn công cụ theo thứ tự các thao tác

Trước khi chọn cấu trúc, hãy hỏi dữ liệu có thay đổi giữa các câu hỏi hay không. Cùng một dãy số nhưng lịch sử thao tác khác nhau có thể cần cách xử lý khác nhau.

Bảng không có nghĩa Fenwick phù hợp với mọi loại truy vấn. Phép trừ hai tiền tố dựa trên việc phép cộng có phép ngược là phép trừ. Chẳng hạn, không thể lấy số lớn nhất của một đoạn bằng cách trừ hai số lớn nhất của hai tiền tố.

a = [9, 2, 3]
b = [9, 7, 3]
print(max(a[:3]), max(a[:1]), max(a[1:3]))
print(max(b[:3]), max(b[:1]), max(b[1:3]))

9 9 3
9 9 7

Hai cặp cực đại tiền tố giống nhau, nhưng cực đại của đoạn còn lại khác nhau. Chỉ biết hai cực đại tiền tố không đủ để suy ra cực đại đoạn. Nếu đề đổi từ tổng sang cực đại, bạn cần một cấu trúc khác.

Lỗi thường gặp

Tóm tắt

• Tiền tố trả lời tổng đoạn tĩnh bằng hiệu hai tổng.
• Mảng hiệu xử lý nhiều phép cộng đoạn trước khi khôi phục dãy.
• Tần suất tiền tố giúp đếm đoạn có tổng xác định, kể cả khi có số âm.
• Fenwick phù hợp khi cập nhật điểm và truy vấn tổng xen kẽ.

Kiểm tra nhanh

Đọc thêm

• Python: cấu trúc dữ liệu từ điển.
• Swiss Olympiad in Informatics: cây Fenwick.