Cấu trúc dữ liệu hỗ trợ truy vấn tổng đoạn (Range Sum) trong Python

Xây dựng cấu trúc dữ liệu được khởi tạo với một mảng số nguyên, hỗ trợ phương thức get_sum(i, j) trả về tổng các phần tử từ chỉ số i đến j-1 (tức là đoạn [i, j)) một cách hiệu quả.

Ví dụ: mảng [5,2,3,6,4,7,8,9,3,2], gọi get_sum(1,5) → 15 (tổng [2,3,6,4]), get_sum(4,8) → 28 (tổng [4,7,8,9]).

Thuật toán (Prefix Sum / Tổng tiền tố)

1. Khởi tạo mảng sums với phần tử đầu tiên là 0.
2. Duyệt mảng đầu vào, mỗi phần tử x thêm x + sums[-1] vào sums. Khi đó sums[k] là tổng k phần tử đầu tiên.
3. get_sum(i, j) trả về sums[j] - sums[i] (tổng từ chỉ số i đến j-1).

Độ phức tạp: khởi tạo O(n), truy vấn O(1).

Mã Python

class RangeSum:
    def __init__(self, array):
        self.sums = [0]
        for x in array:
            self.sums.append(x + self.sums[-1])

    def get_sum(self, i, j):
        return self.sums[j] - self.sums[i]

array = [5, 2, 3, 6, 4, 7, 8, 9, 3, 2]
obj = RangeSum(array)
print(obj.get_sum(1, 5))  # 15
print(obj.get_sum(4, 8))  # 28

Đầu vào / Đầu ra

Input: [5,2,3,6,4,7,8,9,3,2] + các lệnh gọi get_sum

Output: 15 28