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
|