Chương trình tìm số cách tạo mảng công bằng trong Python

Cho mảng nums. Xóa đúng 1 phần tử. Mảng công bằng (fair) nếu tổng chỉ số chẵn = tổng chỉ số lẻ (sau khi xóa, chỉ số dịch chuyển). Đếm số chỉ số xóa làm mảng công bằng.

Ví dụ: nums=[5,3,7,2] → 1. Xóa chỉ số 1: [5,7,2], even=5+2=7, odd=7 → fair.

Giải thuật (Prefix sums chẵn/lẻ)

1. Tính sm1 = tổng chỉ số lẻ (1,3,5...), sm2 = tổng chỉ số chẵn (2,4,6...) của mảng gốc (bỏ index 0).
2. Nếu sm1 == sm2: xóa index 0 làm công bằng → res++.
3. Duyệt i từ 1..n-1, mô phỏng xóa i:
   • Nếu i lẻ: sm1 = sm1 - nums[i] + nums[i-1] (index i-1 chuyển từ chẵn sang lẻ).
   • Nếu i chẵn: sm2 = sm2 - nums[i] + nums[i-1].
   • Nếu sm1 == sm2: res++.

Ví dụ minh họa

def solve(nums):
    res, sm1, sm2 = 0, 0, 0
    for i in range(1, len(nums)):
        if i % 2 == 1:
            sm1 += nums[i]
        else:
            sm2 += nums[i]
    if sm1 == sm2:
        res += 1
    for i in range(1, len(nums)):
        if i % 2 == 1:
            sm1 = sm1 - nums[i] + nums[i-1]
        else:
            sm2 = sm2 - nums[i] + nums[i-1]
        if sm1 == sm2:
            res += 1
    return res

nums = [5,3,7,2]
print(solve(nums))

Đầu vào

[5,3,7,2]

Đầu ra

1