Chương trình đếm số dãy con cấp số cộng độ dài ≥ 3 trong Python

Cho danh sách số nums. Đếm số dãy con (subsequence, không nhất thiết liên tiếp) độ dài ≥ 3 là cấp số cộng (chênh lệch giữa các phần tử liên tiếp bằng nhau).

Ví dụ

nums = [6, 12, 13, 8, 10, 14]

Các dãy con cấp số cộng độ dài ≥ 3:

• [6, 8, 10] (diff = 2)
• [6, 10, 14] (diff = 4)
• [12, 13, 14] (diff = 1)

Kết quả: 3.

Thuật toán (DP với Dictionary)

1. dp[(i, diff)] = số dãy con cấp số cộng kết thúc tại chỉ số i với công sai diff (độ dài ≥ 2).
2. Duyệt i từ 0 đến n-1, với mỗi j < i:
   • diff = nums[i] - nums[j]
   • prev = dp.get((i, diff), 0)
   • prevprev = dp.get((j, diff), 0)
   • dp[(i, diff)] = prev + prevprev + 1
   • res += prevprev (chỉ đếm dãy độ dài ≥ 3)

Triển khai Python

class Solution:
    def solve(self, nums):
        dp = {}
        n = len(nums)
        res = 0
        for i in range(n):
            for j in range(i):
                diff = nums[i] - nums[j]
                prev = dp.get((i, diff), 0)
                prevprev = dp.get((j, diff), 0)
                dp[(i, diff)] = prev + prevprev + 1
                res += prevprev
        return res

# Test
ob = Solution()
print(ob.solve([6, 12, 13, 8, 10, 14]))  # Output: 3

Độ phức tạp

• Thời gian: O(n²).
• Không gian: O(n²) cho dictionary dp.

Bài toán "Arithmetic Slices II - Subsequence" (LeetCode 446).