Mảng cộng dồn (Prefix Sum) và mảng hiệu: giải thích dễ hiểu, code C++, Python

Bài toán mở đầu: Vì sao cộng bằng vòng lặp bị quá thời gian?

Giả sử bạn có một mảng a gồm n = 100.000 phần tử ghi doanh thu mỗi ngày của cửa hàng, và cần trả lời q = 100.000 câu hỏi dạng: "Tổng doanh thu từ ngày thứ l đến ngày thứ r là bao nhiêu?".

Nếu với mỗi câu hỏi bạn đều chạy vòng lặp for từ l đến r để cộng từng số:

• Một câu hỏi hỏi từ đầu mảng đến cuối mảng tốn n phép cộng.
• Thực hiện q câu hỏi như vậy sẽ tốn khoảng n × q = 10¹⁰ phép tính. Như đã phân tích trong bài Độ phức tạp thuật toán Big-O, 10 tỷ phép tính mất hàng chục giây và chắc chắn bị báo lỗi chạy quá thời gian (TLE).

Nguyên nhân gây lãng phí là những đoạn giao nhau bị cộng đi cộng lại hàng chục nghìn lần. Mảng cộng dồn (Prefix Sum) giải quyết triệt để việc này bằng cách cộng dồn một lần duy nhất ngay từ đầu.

Ý tưởng và công thức chuẩn không sợ lệch chỉ số

Xét mảng a có n phần tử đánh chỉ số từ 0 đến n - 1. Thay vì tạo mảng cộng dồn có cùng độ dài n, cách viết sạch nhất và ít lỗi nhất là tạo mảng pref có độ dài n + 1, trong đó:

• pref[0] = 0 (tổng của 0 phần tử đầu tiên).
• pref[k] là tổng của k phần tử đầu tiên trong mảng a (tức là a[0] + a[1] + ... + a[k - 1]).
• Công thức xây dựng nối tiếp: pref[i + 1] = pref[i] + a[i].

Khi cần tính tổng các phần tử từ vị trí l đến vị trí r (tức a[l] + a[l + 1] + ... + a[r]), ta chỉ việc lấy tổng từ đầu mảng đến hết vị trí r (chính là pref[r + 1]), rồi bớt đi phần dư ở phía trước vị trí l (là a[0] + ... + a[l - 1], tức pref[l]):

Tổng đoạn [l..r] = pref[r + 1] − pref[l]

Hãy xem ví dụ với mảng 7 phần tử a = [3, 1, 4, 1, 5, 9, 2] khi cần lấy tổng đoạn từ l = 2 đến r = 5:

• Tổng đoạn [2, 5] (4 + 1 + 5 + 9): pref[5 + 1] - pref[2] = pref[6] - pref[2] = 23 - 4 = 19.
• Tổng đoạn [0, 3] (3 + 1 + 4 + 1): pref[3 + 1] - pref[0] = pref[4] - pref[0] = 9 - 0 = 9. Nhờ có pref[0] = 0, ngay cả khi l = 0 công thức vẫn đúng tuyệt đối mà không cần rẽ nhánh if.

Code mảng cộng dồn 1 chiều bằng C++ và Python

Lưu ý trong C++, dù các phần tử a[i] vừa trong kiểu int (tới 10⁹), tổng của nhiều phần tử cộng lại rất dễ vượt quá 2 × 10⁹, nên mảng pref luôn phải khai báo kiểu long long:

#include <iostream>
#include <vector>
using namespace std;

// Xây dựng mảng cộng dồn độ dài n + 1 (pref[0] = 0)
vector<long long> taoMangCongDon(const vector<int>& a) {
    int n = (int)a.size();
    vector<long long> pref(n + 1, 0);
    for (int i = 0; i < n; ++i) {
        pref[i + 1] = pref[i] + a[i];
    }
    return pref;
}

// Trả về tổng đoạn [l, r] (chỉ số 0-based) trong O(1)
long long tongDoan(const vector<long long>& pref, int l, int r) {
    return pref[r + 1] - pref[l];
}

int main() {
    vector<int> a = {3, 1, 4, 1, 5, 9, 2};
    vector<long long> pref = taoMangCongDon(a);
    cout << tongDoan(pref, 2, 5) << "\n";  // 19 (4 + 1 + 5 + 9)
    cout << tongDoan(pref, 0, 3) << "\n";  // 9  (3 + 1 + 4 + 1)
    return 0;
}

def tao_mang_cong_don(a):
    """Xây dựng mảng cộng dồn độ dài n + 1 (pref[0] = 0)."""
    pref = [0] * (len(a) + 1)
    for i, x in enumerate(a):
        pref[i + 1] = pref[i] + x
    return pref


def tong_doan(pref, l, r):
    """Trả về tổng đoạn [l, r] (chỉ số 0-based) trong O(1)."""
    return pref[r + 1] - pref[l]


a = [3, 1, 4, 1, 5, 9, 2]
pref = tao_mang_cong_don(a)
print(tong_doan(pref, 2, 5))  # 19
print(tong_doan(pref, 0, 3))  # 9

Bài toán ngược: Cập nhật nhiều đoạn nhanh bằng Mảng hiệu (Difference Array)

Mảng cộng dồn giúp lấy tổng đoạn trong O(1). Nhưng nếu đề bài đổi ngược lại: cho mảng ban đầu toàn số 0, thực hiện q thao tác cộng thêm giá trị v vào mọi phần tử thuộc đoạn [l, r], rồi sau khi thực hiện xong tất cả các thao tác mới hỏi mảng cuối cùng là gì?

Lúc này ta dùng phép toán ngược của mảng cộng dồn, gọi là mảng hiệu (Difference Array) diff:

• Muốn cộng thêm v vào đoạn [l, r], ta chỉ sửa đúng 2 vị trí trong O(1):
  • diff[l] += v (bắt đầu từ vị trí l, giá trị cộng dồn tăng thêm v).
  • diff[r + 1] -= v (từ vị trí r + 1 trở đi, triệt tiêu lượng tăng thêm v đó).
• Sau khi xử lý xong toàn bộ q thao tác cập nhật, ta chỉ cần cộng dồn mảng diff một lần duy nhất từ trái sang phải trong O(n) là thu được giá trị cuối cùng của mọi phần tử!

#include <iostream>
#include <vector>
using namespace std;

int main() {
    int n = 5;
    vector<long long> diff(n + 1, 0);

    // Cộng một giá trị v vào mọi phần tử thuộc đoạn [l, r] trong O(1)
    auto congDoan = [&](int l, int r, int v) {
        diff[l] += v;
        diff[r + 1] -= v;
    };

    congDoan(1, 3, 5);   // [0, 5, 5, 5, 0]
    congDoan(2, 4, 10);  // [0, 5, 15, 15, 10]
    congDoan(0, 2, 2);   // [2, 7, 17, 15, 10]

    // Cộng dồn mảng diff một lần để khôi phục mảng kết quả
    vector<long long> a(n);
    long long cur = 0;
    for (int i = 0; i < n; ++i) {
        cur += diff[i];
        a[i] = cur;
        cout << a[i] << (i + 1 == n ? "\n" : " ");  // 2 7 17 15 10
    }
    return 0;
}

n = 5
diff = [0] * (n + 1)


def cong_doan(l, r, v):
    diff[l] += v
    diff[r + 1] -= v


cong_doan(1, 3, 5)   # [0, 5, 5, 5, 0]
cong_doan(2, 4, 10)  # [0, 5, 15, 15, 10]
cong_doan(0, 2, 2)   # [2, 7, 17, 15, 10]

a = []
cur = 0
for i in range(n):
    cur += diff[i]
    a.append(cur)

print(*a)  # 2 7 17 15 10

Mở rộng sang Mảng cộng dồn 2 chiều (trên bảng lưới)

Khi dữ liệu là một bảng hình chữ nhật m hàng, n cột, ta định nghĩa pref[i][j] là tổng tất cả các ô nằm trong góc hình chữ nhật từ ô đầu bảng đến hết hàng i, cột j. Theo nguyên lý bao hàm — loại trừ:

• Khi xây dựng bảng: pref[i][j] = pref[i - 1][j] + pref[i][j - 1] - pref[i - 1][j - 1] + a[i - 1][j - 1] (vì phần góc trên bên trái pref[i - 1][j - 1] được cộng hai lần nên phải trừ bớt một lần).
• Khi tính tổng hình chữ nhật con có góc trên trái (r1, c1) và góc dưới phải (r2, c2) (0-based): Tổng = pref[r2 + 1][c2 + 1] - pref[r1][c2 + 1] - pref[r2 + 1][c1] + pref[r1][c1].

Mọi truy vấn tổng hình chữ nhật con có kích thước bất kỳ đều chỉ tốn đúng 3 phép cộng trừ trong O(1).

Khi nào dùng Mảng cộng dồn, khi nào cần cấu trúc dữ liệu khác?

Lỗi hay gặp

Luyện ngay trên Big-O Coding

Các bài tập dưới đây trên hệ thống Big-O Coding sử dụng trực tiếp mảng cộng dồn 1 chiều, 2 chiều và mảng hiệu:

Câu hỏi thường gặp

Nguồn tham khảo

• USACO Guide: Introduction to Prefix Sums
• cp-algorithms: Search the subsegment with the maximum/minimum sum
• cppreference: std::partial_sum