Fenwick Tree: Cấu trúc dữ liệu tôi ước mình được học sớm hơn

Trong quá trình học thuật toán, có những cấu trúc dữ liệu tôi chủ động tìm hiểu từ sớm vì nghe tên quá nhiều. Nhưng cũng có những cấu trúc dữ liệu mà tôi chỉ thực sự học khi bị một bài toán “ép” phải học.

Fenwick Tree thuộc nhóm thứ hai.

Lần đầu tiên nhìn thấy Fenwick Tree, tôi không hề có ấn tượng gì đặc biệt. Thậm chí tôi còn nghĩ rằng một cấu trúc dữ liệu chỉ gồm vài chục dòng code chắc cũng không có gì quá phức tạp hay quan trọng.

Nhưng càng giải nhiều bài toán hơn, tôi càng nhận ra Fenwick Tree xuất hiện ở rất nhiều nơi. Từ những bài đếm nghịch thế kinh điển, các bài toán thống kê tần suất, cho tới những bài yêu cầu xử lý hàng trăm nghìn truy vấn trong thời gian rất ngắn.

Điều thú vị là Fenwick Tree không phải là một cấu trúc dữ liệu thường được dạy trong những môn nhập môn ở trường đại học. Chúng ta thường học mảng, danh sách liên kết, hàng đợi, ngăn xếp, cây nhị phân hay bảng băm. Trong khi đó, Fenwick Tree thường chỉ bắt đầu xuất hiện khi tìm hiểu các thuật toán nâng cao hoặc tham gia competitive programming.

Trong bài viết này, tôi muốn chia sẻ cách tôi tiếp cận Fenwick Tree thông qua một bài toán thực tế, quá trình từ một lời giải chậm tới một lời giải đủ nhanh, cũng như trực giác giúp tôi hiểu vì sao cấu trúc dữ liệu này lại hoạt động được.

Tôi biết đến Fenwick Tree như thế nào?

Tôi biết đến Fenwick Tree khi giải một bài toán trên Codeforces.

Như thường lệ, việc đầu tiên tôi làm không phải là nghĩ đến cấu trúc dữ liệu nào cả. Tôi luôn cố gắng tìm ra một lời giải đơn giản nhất trước, miễn là nó đúng.

Bài toán

Đề bài đầy đủ có thể xem tại đây. Tóm tắt ngắn gọn như sau.

Cho một mảng a, trong đó a[i] biểu diễn số lượng hộp ở cột thứ i. Ta thực hiện thao tác trượt tất cả các hộp sang phải nhiều nhất có thể và cần tính tổng quãng đường mà các hộp đã di chuyển. Ngoài ra, được phép loại bỏ đúng một hộp nằm trên cùng của bất kỳ cột nào trước khi thực hiện việc trượt. Mục tiêu là tìm cách loại bỏ một hộp sao cho tổng quãng đường di chuyển của các hộp là lớn nhất.

Khi mới đọc đề bài, tôi chưa nghĩ đến tối ưu. Tôi chỉ nghĩ: Trước hết phải tìm được lời giải đúng đã.

Sau khi phân tích, tôi nhận ra bài toán có thể tách thành hai phần tương đối độc lập:

• Tính tổng quãng đường tất cả các hộp di chuyển.
• Tìm xem việc loại bỏ hộp ở cột nào mang lại lợi ích lớn nhất.

Lời giải đầu tiên

Đầu tiên, tôi muốn tính tổng quãng đường tất cả các hộp di chuyển. Nếu xét một hộp nằm ở cột i, quãng đường tối đa nó có thể di chuyển là:

n - 1 - i

Tuy nhiên, nếu ở cùng độ cao đó đã có một số hộp nằm bên phải, quãng đường thực tế sẽ giảm tương ứng. Từ suy nghĩ đó, tôi viết được đoạn code sau:

int64_t ans = 0;
vector<int> cnt(ranges::max(a) + 1, 0);

for (int i = n - 1; i >= 0; i--) {
    for (int h = 1; h <= a[i]; h++) {
        ans += n - cnt[h] - 1 - i;
        cnt[h]++;
    }
}

Sau đó tôi tiếp tục xử lý phần thứ hai.

Nếu bỏ hộp trên cùng của cột i, thì tất cả các cột phía trước có chiều cao không nhỏ hơn a[i] sẽ được hưởng lợi thêm một đơn vị khoảng cách. Tôi cài đặt trực tiếp:

int64_t d = 0;

for (int i = 0; i < n; i++) {
    int64_t cur = 0;

    for (int j = 0; j < i; j++) {
        if (a[j] >= a[i]) cur++;
    }

    d = max(d, cur);
}

Cuối cùng:

cout << ans + d << '\n';

Mọi thứ có vẻ ổn.

Cho đến khi tôi nhìn lại độ phức tạp.

Vấn đề của lời giải

Thoạt nhìn, đoạn code trên không quá dài. Nhưng nếu phân tích kỹ, ta sẽ thấy nhiều vòng lặp lồng nhau. Độ phức tạp xấp xỉ: O(n² + n·maxA).

Trong khi đề bài cho n = 200000. Điều đó đồng nghĩa với hàng chục tỷ phép tính. Nói cách khác: “Chắc chắn TLE”.

Đây là một tình huống tôi gặp rất nhiều trong competitive programming. Tìm ra lời giải đúng không quá khó. Tìm ra lời giải đủ nhanh mới là phần khó.

Nhìn lại bài toán từ góc độ khác

Khi gặp một lời giải O(n²), tôi thường làm một việc khá đơn giản. Thay vì nhìn vào code, tôi nhìn vào công thức mà code đang tính. Sau một lúc biến đổi, tôi nhận ra phần tính tổng quãng đường thực chất có thể viết lại thành:

int64_t ans = 0;

for (int i = 0; i < n; i++) {
    for (int j = 0; j < i; j++) {
        ans += max(0, a[j] - a[i]);
    }
}

Độ phức tạp vẫn là O(n²), nhưng công thức lúc này dễ quan sát hơn.

Với mỗi phần tử a[i], tôi đang cần tính Σ(a[j] - a[i])với điều kiện:

j < i
a[j] > a[i]

Nhìn kỹ hơn một chút:

Σ(a[j] - a[i]) = Σa[j] - a[i] × count

Đây là khoảnh khắc khiến tôi nhìn thấy hướng tối ưu. Để tính đóng góp của a[i], thực ra tôi chỉ cần biết:

• Có bao nhiêu phần tử trước đó lớn hơn a[i]
• Tổng giá trị của các phần tử đó là bao nhiêu

Bài toán lúc này không còn là duyệt mọi cặp phần tử nữa. Nó đã biến thành một bài toán truy vấn động trên dữ liệu đã xuất hiện trước đó. Và đó chính là lúc Fenwick Tree bước lên sân khấu.

Fenwick Tree là gì?

Lần đầu đọc định nghĩa Fenwick Tree, tôi không cảm thấy nó đặc biệt dễ hiểu. Thông thường người ta sẽ mô tả như sau:

Fenwick Tree là một cấu trúc dữ liệu hỗ trợ cập nhật một phần tử và truy vấn tổng prefix trong O(log n).

Định nghĩa này hoàn toàn đúng. Nhưng nếu mới học, tôi nghĩ nó hơi khô khan. Sau này tôi thường tự giải thích đơn giản hơn:

Fenwick Tree là một cách lưu nhiều tổng trung gian khác nhau để tránh phải cộng lại từ đầu mỗi lần truy vấn.

Nếu nhìn theo góc độ đó, Fenwick Tree giống như một phiên bản thông minh hơn của việc cộng dồn thông thường. Thay vì lưu toàn bộ dữ liệu thô, nó lưu sẵn những “mảnh ghép” đủ lớn để ghép lại thành kết quả cần tìm.

Trực giác giúp tôi hiểu Fenwick Tree

Điều khiến tôi bối rối nhất khi mới học là: Tại sao vài dòng code ngắn như vậy lại có thể chạy trong O(log n)?. Sau một thời gian, tôi nhận ra rằng mình đang nghĩ sai về cách dữ liệu được lưu.

Giả sử có mảng:

1 2 3 4 5 6 7 8

Trực giác đầu tiên thường là:

bit[1] lưu phần tử 1
bit[2] lưu phần tử 2
...

Nhưng Fenwick Tree không hoạt động theo cách đó. Nó lưu các đoạn:

[1]
[1..2]
[3]
[1..4]
[5]
[5..6]
[7]
[1..8]

Tức là rất nhiều tổng trung gian ở các kích thước khác nhau. Khi cần tính tổng từ 1 tới 7, thay vì cộng bảy phần tử riêng lẻ, tôi chỉ việc lấy:

[1..4]
[5..6]
[7]

rồi cộng lại.

Đó chính là ý tưởng cốt lõi của Fenwick Tree.

Cài đặt Fenwick Tree

Điều tôi thích nhất ở Fenwick Tree là phần cài đặt cực kỳ ngắn.

Cập nhật

void add(int p, int v) {
    for (; p <= n; p += p & -p) {
        bit[p] += v;
    }
}

Prefix Sum

int64_t sum(int p) {
    int64_t res = 0;

    for (; p > 0; p -= p & -p) {
        res += bit[p];
    }

    return res;
}

Nếu chỉ nhìn code, rất nhiều người sẽ thấy biểu thức này khá bí ẩn:

p & -p

Tôi cũng từng như vậy. Sau này tôi mới biết đây là cách lấy bit 1 thấp nhất của một số, thường được gọi là lowbit.

Ví dụ:

6 = 110
lowbit = 2

12 = 1100
lowbit = 4

Fenwick Tree sử dụng lowbit để nhảy giữa các đoạn dữ liệu có liên quan. Nhờ đó số bước xử lý luôn chỉ vào khoảng O(log n).

Một cách nhớ Fenwick Tree rất dễ

Nếu phải tóm tắt Fenwick Tree bằng một câu, tôi thường nhớ như sau.

Khi cập nhật:

add(pos, val)

ta đi lên các đoạn lớn hơn:

pos += lowbit(pos);

để thông báo rằng vị trí này vừa thay đổi.

Ngược lại, khi truy vấn:

sum(pos)

ta đi xuống:

pos -= lowbit(pos);

để ghép nhiều đoạn nhỏ thành prefix sum cần tìm.

Chỉ cần nhớ hai hướng di chuyển đó, tôi gần như không bao giờ phải học thuộc code Fenwick Tree.

Coordinate Compression

Một điều rất hay là Fenwick Tree không nhất thiết phải hoạt động trên giá trị gốc.

Ví dụ:

a = {100, 1000000, 7};

Tôi không muốn tạo một Fenwick Tree kích thước một triệu chỉ để lưu ba số. Lúc này ta thực hiện coordinate compression.

Sau khi nén:

7       -> 1
100     -> 2
1000000 -> 3

Mọi thao tác sau đó chỉ làm việc trên:

1..3

Đây là một kỹ thuật mà tôi gần như luôn sử dụng cùng Fenwick Tree.

Áp dụng vào bài toán

Sau khi đã hiểu Fenwick Tree hoạt động như thế nào, hãy quay lại bài toán ban đầu. Khi đang duyệt phần tử:

x = a[i]

tôi cần biết hai thông tin về các phần tử đã xuất hiện trước đó:

• Có bao nhiêu phần tử lớn hơn x
• Tổng các phần tử lớn hơn x

Nếu gọi:

count_gt

là số lượng phần tử lớn hơn x, và:

sum_gt

là tổng giá trị của các phần tử đó, thì đóng góp của x vào đáp án sẽ là:

sum_gt - count_gt * x

Nhờ vậy, thay vì duyệt toàn bộ các phần tử trước đó, tôi chỉ cần biết hai giá trị tổng hợp.

Một Fenwick Tree để đếm

Đầu tiên, tôi tạo một Fenwick Tree dùng để lưu số lượng phần tử đã xuất hiện.

fenwick bit_cnt(m);

Khi gặp một giá trị mới:

bit_cnt.add(id, 1);

Nếu muốn biết có bao nhiêu số lớn hơn x:

cnt_gt = bit_cnt.sum(id + 1, m);

Một Fenwick Tree để lưu tổng

Tiếp theo, tôi tạo thêm một Fenwick Tree thứ hai.

fenwick bit_sum(m);

Lần này thay vì cộng 1, tôi cộng x vào cây.

bit_sum.add(id, x);

Khi đó tổng các giá trị lớn hơn x là:

sum_gt = bit_sum.sum(id + 1, m);

Tính phần tăng thêm khi loại bỏ một hộp

Điều khá thú vị là phần thứ hai của bài toán cũng có thể giải bằng chính Fenwick Tree.

Trong lời giải ban đầu, tôi viết:

for (int j = 0; j < i; j++) {
    if (a[j] >= a[i]) {
        cur++;
    }
}

Đoạn code này đang trả lời câu hỏi: Có bao nhiêu phần tử trước đó >= a[i]? Mà đây lại chính là loại truy vấn Fenwick Tree xử lý rất tốt.

Nếu đang đứng tại:

x = a[i]

thì:

bit_cnt.sum(id, m)

chính là số lượng phần tử đã xuất hiện có giá trị lớn hơn hoặc bằng x.

Từ đó ta có thể cập nhật trực tiếp:

d = max(d, bit_cnt.sum(id, m));

mà không cần thêm một vòng lặp O(n) nào nữa. Tất cả được thực hiện trong một lần duyệt Điều tôi thích nhất ở lời giải này là mọi thứ đều diễn ra trong cùng một vòng lặp.

Khi duyệt một phần tử:

x = a[i]

ta thực hiện:

• Truy vấn số lượng phần tử lớn hơn x
• Truy vấn tổng các phần tử lớn hơn x
• Cập nhật đáp án cho phần thứ nhất
• Cập nhật đáp án cho phần thứ hai
• Thêm x vào hai Fenwick Tree

Tất cả các thao tác đều có độ phức tạp O(log n), nên toàn bộ lời giải trở thành O(n log n).

Code hoàn chỉnh

#include <bits/stdc++.h>
using namespace std;

struct fenwick {
    int n;
    vector<int64_t> bit;

    fenwick(int n_) : n(n_), bit(n + 1, 0) {}

    void add(int p, int64_t v) {
        for (; p <= n; p += p & -p) {
            bit[p] += v;
        }
    }

    int64_t sum(int p) const {
        int64_t res = 0;

        for (; p > 0; p -= p & -p) {
            res += bit[p];
        }

        return res;
    }

    int64_t sum(int l, int r) const {
        return sum(r) - sum(l - 1);
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int t;
    cin >> t;

    while (t--) {
        int n;
        cin >> n;

        vector<int> a(n);

        for (int &x : a) cin >> x;

        vector<int> vals = a;

        ranges::sort(vals);
        vals.erase(unique(vals.begin(), vals.end()), vals.end());

        int m = (int)vals.size();

        auto gid = int x {
            return int(lower_bound(vals.begin(), vals.end(), x) - vals.begin()) + 1;
        };

        fenwick bit_cnt(m);
        fenwick bit_sum(m);

        int64_t ans = 0;
        int64_t d = 0;

        for (int x : a) {
            int id = gid(x);

            int64_t cnt_gt = bit_cnt.sum(id + 1, m),
                sum_gt = bit_sum.sum(id + 1, m);

            ans += sum_gt - cnt_gt * x;
            d = max(d, bit_cnt.sum(id, m));

            bit_cnt.add(id, 1);
            bit_sum.add(id, x);
        }

        cout << ans + d << '\n';
    }

    return 0;
}

Đây là một trong những ví dụ khiến tôi thực sự hiểu giá trị của Fenwick Tree. Nếu chỉ nhìn vào lời giải cuối cùng, rất dễ nghĩ rằng Fenwick Tree là phần quan trọng nhất. Nhưng nhìn lại quá trình giải bài, tôi thấy bước quan trọng hơn lại là việc biến đổi công thức và nhận ra mình chỉ cần biết “bao nhiêu phần tử lớn hơn x” và “tổng của chúng là bao nhiêu”.

Fenwick Tree chỉ là công cụ giúp trả lời hai câu hỏi đó một cách hiệu quả. Đây cũng là điều tôi gặp lại rất nhiều lần khi giải các bài toán khác: tối ưu thường bắt đầu từ việc nhìn lại bài toán dưới một góc độ khác, trước khi chọn cấu trúc dữ liệu phù hợp.

Những trường hợp tôi thường nghĩ đến Fenwick Tree

Sau một thời gian giải thuật toán, tôi nhận ra Fenwick Tree xuất hiện rất thường xuyên trong một số nhóm bài.

Ví dụ:

• Đếm nghịch thế (Inversion Count)
• Đếm số lượng phần tử nhỏ hơn hoặc lớn hơn một giá trị
• Truy vấn tần suất xuất hiện
• Xếp hạng (Ranking)
• Order Statistics
• Thống kê động trên dữ liệu đã xuất hiện

Mỗi khi nhìn thấy các cụm từ như “đếm số lượng”, hoặc “tổng các phần tử thuộc một khoảng giá trị”, Fenwick Tree gần như là một trong những công cụ đầu tiên tôi nghĩ tới.

Điều tôi rút ra sau khi học Fenwick Tree

Điều khiến tôi thích Fenwick Tree không phải là vì nó nhanh. Điều khiến tôi thích nó là vì nó thay đổi cách tôi nhìn bài toán. Rất nhiều bài toán ban đầu trông như đang yêu cầu duyệt mọi cặp phần tử, nhưng sau khi biến đổi công thức một chút, chúng lại trở thành đếm số lượng trong một khoảng hoặc tính tổng trong một khoảng.

Khi đó Fenwick Tree xuất hiện như một lời giải rất tự nhiên. Và đó cũng là lúc tôi nhận ra rằng việc học cấu trúc dữ liệu không chỉ để nhớ cách cài đặt. Điều quan trọng hơn là học cách nhận ra bài toán nào phù hợp với cấu trúc dữ liệu đó.

Kết luận

Fenwick Tree là một trong những cấu trúc dữ liệu nhỏ gọn nhất mà tôi từng học.

Toàn bộ phần cài đặt chỉ gồm vài hàm ngắn, nhưng khả năng ứng dụng của nó lại rất rộng. Từ đếm nghịch thế, thống kê tần suất, xếp hạng cho tới rất nhiều bài toán tối ưu trong competitive programming đều có thể xuất hiện Fenwick Tree ở đâu đó trong lời giải.

Đối với tôi, đây cũng là một trong những cấu trúc dữ liệu mang lại cảm giác “khai sáng” rõ rệt nhất. Không phải vì cách cài đặt phức tạp hay kỹ thuật quá cao siêu, mà vì nó cho thấy một điều rất quan trọng:

Đôi khi chìa khóa để tối ưu không nằm ở việc viết code nhanh hơn, mà nằm ở việc nhìn bài toán theo một góc độ khác.

Nếu bạn đang bắt đầu tìm hiểu các kỹ thuật xử lý truy vấn trên mảng và muốn bước xa hơn khỏi những vòng lặp O(n²), tôi nghĩ Fenwick Tree là một trong những cấu trúc dữ liệu đáng học nhất. Và cũng giống như tôi khi mới học nó, rất có thể bạn sẽ bắt đầu nhìn thấy Fenwick Tree xuất hiện ở khắp mọi nơi sau đó.