Chương trình đếm số mảng con đẹp trong Python

Mô tả bài toán

Cho một mảng nums và một giá trị k. Một mảng con được gọi là "đẹp" nếu nó chứa đúng k số lẻ. Hãy tìm số lượng mảng con đẹp.

Ví dụ

Nếu đầu vào là nums = [1,1,2,1,1], k = 3, thì đầu ra sẽ là 2 vì có hai mảng con đẹp: [1,1,2,1] và [1,2,1,1].

Giải thuật

1. Tạo danh sách odd_i lưu chỉ số của các phần tử lẻ trong nums.
2. Khởi tạo start = 0, end = k - 1, i = 0, count = 0.
3. Trong khi end < len(odd_i):
   • Nếu end == len(odd_i) - 1, gán j = len(nums) - 1; ngược lại j = odd_i[end + 1] - 1.
   • Cập nhật count += (odd_i[start] - i + 1) * (j - odd_i[end] + 1).
   • Cập nhật i = odd_i[start] + 1, start += 1, end += 1.
4. Trả về count.

Triển khai Python

def solve(nums, k):
    odd_i = []
    for i in range(len(nums)):
        if nums[i] % 2 == 1:
            odd_i.append(i)
    start = 0
    end = k - 1
    i = 0
    count = 0
    while end < len(odd_i):
        if end == len(odd_i) - 1:
            j = len(nums) - 1
        else:
            j = odd_i[end + 1] - 1
        count += (odd_i[start] - i + 1) * (j - odd_i[end] + 1)
        i = odd_i[start] + 1
        start += 1
        end += 1
    return count

nums = [1, 1, 2, 1, 1]
k = 3
print(solve(nums, k))

Đầu vào

[1,1,2,1,1], 3

Đầu ra

2