[leetcode] 413. Arithmetic Slices

题目

An integer array is called arithmetic if it consists of at least three elements and if the difference between any two consecutive elements is the same.

• For example, [1,3,5,7,9], [7,7,7,7], and [3,-1,-5,-9] are arithmetic sequences.

Given an integer array nums, return the number of arithmetic subarrays of nums.

A subarray is a contiguous subsequence of the array.

Example 1:

Input: nums = [1,2,3,4]
Output: 3
Explanation: We have 3 arithmetic slices in nums: [1, 2, 3], [2, 3, 4] and [1,2,3,4] itself.

Example 2:

Input: nums = [1]
Output: 0

Constraints:

• 1 <= nums.length <= 5000
• -1000 <= nums[i] <= 1000

思路

解法1：首先确认一个等差序列n中三个元素以上等差序列的数量为：1+2+3+...+(n-3)，同样为等差序列，因此输入一个等差序列的长度，可以得出其子等差序列的数量。接着遍历序列，获取所有可能的等差序列即可。

解法2：动态规划，如果nums[i]-nums[i-1] == nums[i-1]-nums[i-2]，那么dp[i] = dp[i-1]+1。

代码

python版本：

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        def count(n):
            return (n-2)*(n-1)//2 if n >= 3 else 0
        if len(nums) < 3:
            return 0
        res = 0
        l = 0
        d = nums[1]-nums[0]
        for i, num in enumerate(nums):
            if i >= 1 and num-nums[i-1] != d:
                res += count(i-l)
                d = num-nums[i-1]
                l = i-1
        if nums[-1]-nums[-2] == d:
            res += count(len(nums)-l)
        return res

# 动态规划
class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        dp = [0]*len(nums)
        res = 0
        for i in range(2, len(nums)):
            if nums[i]-nums[i-1] == nums[i-1]-nums[i-2]:
                dp[i] = dp[i-1]+1
                res += dp[i]
        return res