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
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        if len(nums) < 3:
            return 0
        seq = []
        count = 0
        for i in range(2, len(nums)):
            if nums[i] - nums[i-1] == nums[i-1] - nums[i-2]:
                count += 1
            elif count > 0:
                seq.append(count + 2)
                count = 0
        if count > 0:
            seq.append(count + 2)
        res = 0
        for s in seq:
            # 1 + 2 + ... + (n-2) = (1 + (n-2)) * (n-2) / 2
            res += (s*s - s*3 + 2) // 2
        return res


'''
Simplier math solution
'''
class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        count = 0
        total = 0
        for i in range(2, len(nums)):
            if nums[i] - nums[i-1] == nums[i-1] - nums[i-2]:
                count += 1
            else:
                total += (count + 1) * (count) // 2
                count = 0
        total += count * (count + 1) // 2
        return total


'''
DP:
'''
class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        dp = 0
        total = 0
        for i in range(2, len(nums)):
            if nums[i] - nums[i-1] == nums[i-1] - nums[i-2]:
                dp += 1
                total += dp
            else:
                dp = 0
        return total