Bài 7.1 - Kỹ thuật Prefix Sum trong Python

Bài 7.1. Kỹ thuật Prefix Sum trong Python
Chào mừng bạn đến với bài học về một kỹ thuật cực kỳ hữu ích trong lập trình, đặc biệt là khi làm việc với mảng và các bài toán liên quan đến tính tổng trên các đoạn con: Kỹ thuật Prefix Sum (Tổng tiền tố). Nếu bạn thường xuyên gặp phải các bài toán yêu cầu tính tổng các phần tử trong một đoạn liên tục của mảng một cách nhanh chóng, thì kỹ thuật này chính là người bạn đồng hành không thể thiếu của bạn.
Trong bài viết này, chúng ta sẽ đi sâu vào:
- Prefix Sum là gì? Hiểu rõ bản chất của kỹ thuật này.
- Cách xây dựng mảng Prefix Sum trong Python.
- Ứng dụng chính: Giải quyết bài toán truy vấn tổng đoạn (Range Sum Query) một cách hiệu quả vượt trội.
- Các ví dụ minh họa chi tiết bằng code Python.
Hãy cùng bắt đầu!
Prefix Sum là gì?
Đơn giản mà nói, mảng Prefix Sum (hay mảng tổng tiền tố) của một mảng gốc arr
là một mảng mới, gọi là prefix_sum_arr
, có cùng kích thước (hoặc lớn hơn một chút tùy cách cài đặt) mà ở mỗi chỉ số i
, giá trị của prefix_sum_arr[i]
là tổng của tất cả các phần tử từ đầu mảng gốc (chỉ số 0) đến chỉ số i
(bao gồm cả i
).
Hay nói cách khác:
prefix_sum_arr[i] = arr[0] + arr[1] + ... + arr[i]
Với trường hợp đặc biệt prefix_sum_arr[0] = arr[0]
.
Cách tính này có thể được tối ưu hơn bằng công thức truy hồi:
prefix_sum_arr[i] = prefix_sum_arr[i-1] + arr[i]
(với i > 0
)
Công thức này giúp ta tính giá trị tiếp theo chỉ dựa vào giá trị trước đó và phần tử hiện tại của mảng gốc, tránh việc phải tính lại tổng từ đầu mỗi lần.
Cách Xây dựng Mảng Prefix Sum
Việc xây dựng mảng Prefix Sum rất đơn giản, chỉ cần một lần duyệt qua mảng gốc. Chúng ta có thể sử dụng công thức truy hồi đã nêu ở trên.
Một cách cài đặt phổ biến và khá tiện lợi là thêm một giá trị 0
vào đầu mảng Prefix Sum. Tức là:
prefix_sum_arr[0] = 0
prefix_sum_arr[i] = arr[0] + arr[1] + ... + arr[i-1]
(với i > 0
)
Trong trường hợp này, mảng prefix_sum_arr
sẽ có kích thước lớn hơn mảng arr
một đơn vị. Giá trị prefix_sum_arr[i]
lúc này sẽ là tổng của các phần tử từ arr[0]
đến arr[i-1]
.
Ưu điểm của cách này sẽ được thấy rõ khi chúng ta thực hiện truy vấn tổng đoạn.
Hãy xem ví dụ code để xây dựng mảng Prefix Sum theo cách này.
Ví dụ 1: Xây dựng Mảng Prefix Sum
Giả sử chúng ta có mảng gốc arr = [10, 4, 8, 3, 5]
. Chúng ta muốn xây dựng mảng Prefix Sum của nó.
def build_prefix_sum(arr):
"""
Xây dựng mảng Prefix Sum với giá trị 0 ở đầu.
prefix_sum_arr[i] = tổng các phần tử từ arr[0] đến arr[i-1]
"""
n = len(arr)
prefix_sum_arr = [0] * (n + 1) # Khởi tạo mảng prefix sum có kích thước n+1 với giá trị 0 ở đầu
# Duyệt qua mảng gốc và tính toán các giá trị prefix sum
for i in range(n):
prefix_sum_arr[i+1] = prefix_sum_arr[i] + arr[i]
return prefix_sum_arr
# Mảng gốc
original_array = [10, 4, 8, 3, 5]
# Xây dựng mảng prefix sum
prefix_sums = build_prefix_sum(original_array)
print("Mảng gốc:", original_array)
print("Mảng Prefix Sum (với 0 ở đầu):", prefix_sums)
# Expected output:
# Mảng gốc: [10, 4, 8, 3, 5]
# Mảng Prefix Sum (với 0 ở đầu): [0, 10, 14, 22, 25, 30]
Giải thích Code:
def build_prefix_sum(arr):
: Định nghĩa hàm nhận vào mảng gốcarr
.n = len(arr)
: Lấy độ dài của mảng gốc.prefix_sum_arr = [0] * (n + 1)
: Khởi tạo một mảngprefix_sum_arr
mới có kích thướcn + 1
. Tất cả các phần tử ban đầu được đặt là0
. Phần tử đầu tiênprefix_sum_arr[0]
sẽ luôn là0
theo cách cài đặt này.for i in range(n):
: Vòng lặp duyệt qua các chỉ số của mảng gốcarr
từ0
đếnn-1
.prefix_sum_arr[i+1] = prefix_sum_arr[i] + arr[i]
: Đây là công thức truy hồi.prefix_sum_arr[i]
là tổng của các phần tử từarr[0]
đếnarr[i-1]
.arr[i]
là phần tử hiện tại của mảng gốc.prefix_sum_arr[i+1]
sẽ lưu tổng của các phần tử từarr[0]
đếnarr[i]
(tức là tổng của các phần tử từarr[0]
đếnarr[i-1]
cộng thêmarr[i]
).- Lưu ý chỉ số: vì
prefix_sum_arr
có thêm số0
ở đầu, chỉ sối
củaarr
tương ứng với chỉ sối+1
trongprefix_sum_arr
.
return prefix_sum_arr
: Trả về mảng Prefix Sum đã xây dựng.
Sau khi chạy code, bạn sẽ thấy mảng Prefix Sum [0, 10, 14, 22, 25, 30]
. Hãy kiểm tra lại:
prefix_sum_arr[0]
= 0 (theo cài đặt)prefix_sum_arr[1]
=arr[0]
= 10prefix_sum_arr[2]
=arr[0]
+arr[1]
= 10 + 4 = 14prefix_sum_arr[3]
=arr[0]
+arr[1]
+arr[2]
= 10 + 4 + 8 = 22prefix_sum_arr[4]
=arr[0]
+arr[1]
+arr[2]
+arr[3]
= 10 + 4 + 8 + 3 = 25prefix_sum_arr[5]
=arr[0]
+arr[1]
+arr[2]
+arr[3]
+arr[4]
= 10 + 4 + 8 + 3 + 5 = 30
Mảng Prefix Sum đã được xây dựng chính xác. Thời gian để xây dựng mảng này là O(n), với n
là độ dài của mảng gốc, vì chúng ta chỉ duyệt qua mảng một lần.
Ứng dụng: Truy vấn Tổng Đoạn (Range Sum Query)
Đây là lý do chính khiến kỹ thuật Prefix Sum trở nên mạnh mẽ. Bài toán đặt ra là: tính tổng các phần tử trong mảng gốc từ chỉ số start
đến chỉ số end
(bao gồm cả start
và end
), với 0 <= start <= end < n
.
Cách Truy vấn Tổng Đoạn Naive (Thông thường)
Cách đơn giản nhất là dùng một vòng lặp hoặc hàm sum()
của Python để cộng trực tiếp các phần tử trong đoạn đó.
Ví dụ 2: Truy vấn Tổng Đoạn Cách Naive
def naive_range_sum(arr, start, end):
"""
Tính tổng các phần tử trong đoạn [start, end] của mảng gốc
bằng cách duyệt trực tiếp.
"""
if start < 0 or end >= len(arr) or start > end:
return "Chỉ số không hợp lệ" # Xử lý trường hợp ngoại lệ
total = 0
for i in range(start, end + 1):
total += arr[i]
return total
# Hoặc đơn giản hơn trong Python: return sum(arr[start : end + 1])
original_array = [10, 4, 8, 3, 5]
# Các truy vấn ví dụ
print("Truy vấn naive sum từ 1 đến 3:", naive_range_sum(original_array, 1, 3))
print("Truy vấn naive sum từ 0 đến 2:", naive_range_sum(original_array, 0, 2))
print("Truy vấn naive sum từ 2 đến 4:", naive_range_sum(original_array, 2, 4))
# Expected output:
# Truy vấn naive sum từ 1 đến 3: 15 (4 + 8 + 3)
# Truy vấn naive sum từ 0 đến 2: 22 (10 + 4 + 8)
# Truy vấn naive sum từ 2 đến 4: 16 (8 + 3 + 5)
Giải thích Code:
- Hàm
naive_range_sum
nhận mảngarr
và hai chỉ sốstart
,end
. - Nó kiểm tra các chỉ số đầu vào có hợp lệ không.
- Vòng lặp
for i in range(start, end + 1):
duyệt qua tất cả các phần tử trong đoạn từstart
đếnend
. total += arr[i]
cộng dồn giá trị của từng phần tử vào biếntotal
.- Thời gian thực hiện của mỗi truy vấn này phụ thuộc vào độ dài của đoạn
[start, end]
. Trong trường hợp xấu nhất (đoạn là cả mảng), thời gian là O(n).
Cách naive này hoạt động tốt nếu bạn chỉ cần thực hiện một vài truy vấn. Tuy nhiên, nếu bạn cần thực hiện rất nhiều truy vấn trên cùng một mảng, tổng thời gian sẽ là O(k n), với k
là số lượng truy vấn. Điều này có thể trở nên rất chậm* khi k
và n
lớn.
Cách Truy vấn Tổng Đoạn bằng Prefix Sum
Đây là lúc sức mạnh của Prefix Sum được thể hiện. Nhờ cách xây dựng mảng Prefix Sum (với số 0 ở đầu), chúng ta có thể tính tổng của đoạn arr[start : end + 1]
(từ chỉ số start
đến end
của mảng gốc) chỉ bằng một phép trừ đơn giản:
sum(arr[start : end + 1]) = prefix_sum_arr[end + 1] - prefix_sum_arr[start]
Tại sao lại như vậy?
prefix_sum_arr[end + 1]
là tổng của các phần tử từarr[0]
đếnarr[end]
.prefix_sum_arr[start]
là tổng của các phần tử từarr[0]
đếnarr[start - 1]
.
Khi ta lấy prefix_sum_arr[end + 1] - prefix_sum_arr[start]
, ta đang lấy (tổng từ 0 đến end
) trừ đi (tổng từ 0 đến start - 1
). Kết quả chính xác là tổng của các phần tử từ arr[start]
đến arr[end]
.
Điều tuyệt vời là phép tính này chỉ mất O(1) thời gian!
Ví dụ 3: Truy vấn Tổng Đoạn bằng Prefix Sum
Chúng ta sẽ sử dụng mảng Prefix Sum đã xây dựng ở Ví dụ 1.
def prefix_sum_range_sum(prefix_sum_arr, start, end):
"""
Tính tổng các phần tử trong đoạn [start, end] của mảng gốc
sử dụng mảng Prefix Sum.
"""
# Kiểm tra chỉ số hợp lệ dựa trên mảng gốc (có độ dài len(prefix_sum_arr) - 1)
n_original = len(prefix_sum_arr) - 1
if start < 0 or end >= n_original or start > end:
return "Chỉ số không hợp lệ"
# Công thức kỳ diệu!
# Tổng từ arr[start] đến arr[end]
# = (Tổng từ arr[0] đến arr[end]) - (Tổng từ arr[0] đến arr[start-1])
# = prefix_sum_arr[end + 1] - prefix_sum_arr[start]
return prefix_sum_arr[end + 1] - prefix_sum_arr[start]
# Sử dụng mảng prefix_sums đã tạo từ original_array = [10, 4, 8, 3, 5]
# prefix_sums = [0, 10, 14, 22, 25, 30]
# Các truy vấn ví dụ (giống với ví dụ naive)
print("Truy vấn prefix sum từ 1 đến 3:", prefix_sum_range_sum(prefix_sums, 1, 3))
print("Truy vấn prefix sum từ 0 đến 2:", prefix_sum_range_sum(prefix_sums, 0, 2))
print("Truy vấn prefix sum từ 2 đến 4:", prefix_sum_range_sum(prefix_sums, 2, 4))
print("Truy vấn prefix sum từ 0 đến 4 (toàn bộ mảng):", prefix_sum_range_sum(prefix_sums, 0, 4))
# Expected output:
# Truy vấn prefix sum từ 1 đến 3: 15
# Truy vấn prefix sum từ 0 đến 2: 22
# Truy vấn prefix sum từ 2 đến 4: 16
# Truy vấn prefix sum từ 0 đến 4 (toàn bộ mảng): 30
Giải thích Code:
- Hàm
prefix_sum_range_sum
nhận mảngprefix_sum_arr
và hai chỉ sốstart
,end
(tương ứng với mảng gốc). - Nó kiểm tra chỉ số hợp lệ, lưu ý rằng chỉ số
end
không được vượt quá độ dài của mảng gốc trừ 1, và chỉ sốstart
không được âm. Độ dài mảng gốc được suy ra từ độ dài mảng prefix sum (len(prefix_sum_arr) - 1
). return prefix_sum_arr[end + 1] - prefix_sum_arr[start]
: Đây là dòng mấu chốt. Nó trực tiếp áp dụng công thứcP[j+1] - P[i]
để tính tổng đoạn từ chỉ sối
đếnj
của mảng gốc. Vớistart
lài
vàend
làj
.- Thời gian thực hiện của mỗi truy vấn này là O(1), vì nó chỉ bao gồm một vài phép tính số học cơ bản và truy cập mảng.
So sánh Hiệu quả
Bây giờ chúng ta hãy nhìn vào bức tranh toàn cảnh. Nếu bạn cần thực hiện k
truy vấn tổng đoạn trên một mảng có n
phần tử:
- Cách Naive: Thời gian mỗi truy vấn là O(n) (trường hợp xấu nhất). Tổng thời gian cho
k
truy vấn là O(k * n). - Kỹ thuật Prefix Sum: Thời gian chuẩn bị (xây dựng mảng Prefix Sum) là O(n). Thời gian mỗi truy vấn là O(1). Tổng thời gian cho
k
truy vấn là O(n + k * 1) = O(n + k).
Khi số lượng truy vấn k
rất lớn so với n
, sự khác biệt về hiệu suất là đáng kinh ngạc. Ví dụ, với n=1000 và k=1000000, cách naive sẽ tốn khoảng 10^9 phép tính (rất chậm), trong khi Prefix Sum chỉ tốn khoảng 10^3 + 10^6 = 10^6 phép tính (nhanh hơn gấp nghìn lần).
Ví dụ Tổng Hợp: Nhiều Truy Vấn
Hãy xem một ví dụ đầy đủ hơn, nơi chúng ta thực hiện nhiều truy vấn trên cùng một mảng để thấy rõ lợi ích.
Ví dụ 4: Nhiều Truy Vấn Tổng Đoạn Hiệu Quả
# Mảng gốc
data_array = [2, -1, 5, 8, -3, 6, 12, -7]
print("Mảng dữ liệu gốc:", data_array)
# 1. Xây dựng mảng Prefix Sum
# prefix_sum_arr[i] = tổng các phần tử từ arr[0] đến arr[i-1]
def build_prefix_sum_optimized(arr):
n = len(arr)
prefix_sum_arr = [0] * (n + 1)
for i in range(n):
prefix_sum_arr[i+1] = prefix_sum_arr[i] + arr[i]
return prefix_sum_arr
prefix_sums_data = build_prefix_sum_optimized(data_array)
print("Mảng Prefix Sum:", prefix_sums_data)
# Expected prefix_sums_data: [0, 2, 1, 6, 14, 11, 17, 29, 22]
# 2. Hàm truy vấn sử dụng Prefix Sum
def query_range_sum(prefix_sum_arr, start_idx, end_idx):
"""
Tính tổng đoạn [start_idx, end_idx] của mảng gốc
sử dụng mảng prefix_sum_arr.
"""
# Lấy độ dài mảng gốc từ mảng prefix sum
n_original = len(prefix_sum_arr) - 1
# Kiểm tra chỉ số hợp lệ
if start_idx < 0 or end_idx >= n_original or start_idx > end_idx:
return f"Lỗi: Chỉ số đoạn [{start_idx}, {end_idx}] không hợp lệ cho mảng có độ dài {n_original}"
# Tính tổng = prefix_sum_arr[end_idx + 1] - prefix_sum_arr[start_idx]
return prefix_sum_arr[end_idx + 1] - prefix_sum_arr[start_idx]
# 3. Thực hiện nhiều truy vấn
queries = [(0, 2), (3, 6), (1, 5), (4, 4), (0, 7), (5, 7)]
print("\nKết quả các truy vấn:")
for start, end in queries:
sum_result = query_range_sum(prefix_sums_data, start, end)
print(f"Tổng đoạn từ chỉ số {start} đến {end} là: {sum_result}")
# Expected output (verify manually):
# Tổng đoạn từ chỉ số 0 đến 2: 2 + (-1) + 5 = 6
# Tổng đoạn từ chỉ số 3 đến 6: 8 + (-3) + 6 + 12 = 23
# Tổng đoạn từ chỉ số 1 đến 5: (-1) + 5 + 8 + (-3) + 6 = 15
# Tổng đoạn từ chỉ số 4 đến 4: -3
# Tổng đoạn từ chỉ số 0 đến 7: 2 + (-1) + 5 + 8 + (-3) + 6 + 12 + (-7) = 22
# Tổng đoạn từ chỉ số 5 đến 7: 6 + 12 + (-7) = 11
Giải thích Code:
- Chúng ta bắt đầu với
data_array
. - Hàm
build_prefix_sum_optimized
(tương tự Ví dụ 1) được sử dụng để tạoprefix_sums_data
. Việc này chỉ mất O(n) thời gian một lần duy nhất. - Hàm
query_range_sum
(tương tự Ví dụ 3) nhận mảngprefix_sums_data
và chỉ sốstart_idx
,end_idx
của đoạn cần tính tổng trong mảng gốc. - Chúng ta định nghĩa một danh sách
queries
chứa các cặp chỉ số(start, end)
cho các đoạn cần tính tổng. - Vòng lặp duyệt qua từng truy vấn trong
queries
. Đối với mỗi truy vấn, hàmquery_range_sum
được gọi. - Nhờ có mảng Prefix Sum, mỗi lần gọi
query_range_sum
chỉ mất O(1) thời gian. Tổng thời gian cho tất cả các truy vấn chỉ là O(k), cộng với O(n) thời gian chuẩn bị ban đầu.
Ví dụ này minh họa rõ ràng cách kỹ thuật Prefix Sum chia bài toán thành hai giai đoạn:
- Giai đoạn Tiền xử lý (Precomputation): Xây dựng mảng Prefix Sum (O(n) thời gian).
- Giai đoạn Truy vấn (Query): Thực hiện các truy vấn tổng đoạn (O(1) thời gian mỗi truy vấn).
Nếu số lượng truy vấn lớn, tổng thời gian O(n + k) sẽ vượt trội hoàn toàn so với O(k * n) của phương pháp naive.
Kết luận Tạm thời
Kỹ thuật Prefix Sum là một công cụ mạnh mẽ để giải quyết bài toán truy vấn tổng đoạn trên mảng một chiều một cách hiệu quả. Bằng cách thực hiện một bước tiền xử lý tốn O(n) thời gian để xây dựng mảng tổng tiền tố, chúng ta có thể trả lời bất kỳ truy vấn tổng đoạn nào chỉ trong O(1) thời gian. Điều này làm cho nó trở nên cực kỳ hữu ích trong các tình huống cần thực hiện nhiều truy vấn trên cùng một tập dữ liệu.
Hãy ghi nhớ kỹ thuật này, nó xuất hiện khá thường xuyên trong các bài toán lập trình thi đấu và các ứng dụng thực tế cần hiệu suất cao!
Comments