Chương trình Python đếm số phần tử có thể xóa để tổng chỉ số chẵn/lẻ bằng nhau

Cho danh sách nums. Hàm f(i) xóa phần tử tại chỉ số i và trả về True nếu mảng còn lại có tổng các phần tử ở chỉ số chẵn bằng tổng ở chỉ số lẻ. Tìm số chỉ số i thỏa mãn điều kiện này.

Ví dụ

Input: nums = [6, 8, 5, 2, 3]

Output: 2

Giải thích:

• Xóa 8 (chỉ số 1) → [6, 5, 2, 3]: chỉ số chẵn (0,2): 6+2=8; chỉ số lẻ (1,3): 5+3=8.
• Xóa 2 (chỉ số 3) → [6, 8, 5, 3]: chẵn: 6+5=11; lẻ: 8+3=11.

Thuật toán (Prefix Sum)

1. n = len(nums). Tạo mảng 2D a[2][n+1] lưu prefix sum cho chỉ số chẵn (hàng 0) và lẻ (hàng 1).
2. Duyệt nums, cập nhật prefix sum: a[i%2][i+1] += nums[i].
3. s = sum(nums). Duyệt từng chỉ số i:
   • Tính tổng chỉ số chẵn sau khi xóa i: e = a[0][i] + (a[1][n] - a[1][i+1]).
   • Nếu 2*e == s - nums[i] thì đếm++.

Mã nguồn Python

def solve(nums):
    n = len(nums)
    a = [[0] * (n + 1), [0] * (n + 1)]
    for i, x in enumerate(nums):
        a[0][i + 1] = a[0][i]
        a[1][i + 1] = a[1][i]
        a[i % 2][i + 1] += x

    c = 0
    s = sum(nums)
    for i in range(n):
        e = a[0][i] + a[1][n] - a[1][i + 1]
        if e * 2 == s - nums[i]:
            c += 1
    return c

# Test
nums = [6, 8, 5, 2, 3]
print(solve(nums))  # Output: 2

Độ phức tạp

• Thời gian: O(n)
• Không gian: O(n)