413. Arithmetic Slices 

Description

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

Solutions

Solution 1: Iteration and Counting

We use $d$ to represent the current difference between two adjacent elements, and $cnt$ to represent the length of the current arithmetic sequence. Initially, $d = 3000$, $cnt = 2$.

We iterate through the array nums. For two adjacent elements $a$ and $b$, if $b - a = d$, it means that the current element $b$ also belongs to the current arithmetic sequence, and we increment $cnt$ by 1. Otherwise, it means that the current element $b$ does not belong to the current arithmetic sequence, and we update $d = b - a$, and $cnt = 2$. If $cnt \ge 3$, it means that the length of the current arithmetic sequence is at least 3, and the number of arithmetic sequences is $cnt - 2$, which we add to the answer.

After the iteration, we can get the answer.

In the code implementation, we can also initialize $cnt$ to $0$, and when resetting $cnt$, we directly set $cnt$ to $0$. When adding to the answer, we directly add $cnt$.

The time complexity is $O(n)$, and the space complexity is $O(1)$. Where $n$ is the length of the array nums.

Python3

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        ans = cnt = 0
        d = 3000
        for a, b in pairwise(nums):
            if b - a == d:
                cnt += 1
            else:
                d = b - a
                cnt = 0
            ans += cnt
        return ans

Java

class Solution {
    public int numberOfArithmeticSlices(int[] nums) {
        int ans = 0, cnt = 0;
        int d = 3000;
        for (int i = 0; i < nums.length - 1; ++i) {
            if (nums[i + 1] - nums[i] == d) {
                ++cnt;
            } else {
                d = nums[i + 1] - nums[i];
                cnt = 0;
            }
            ans += cnt;
        }
        return ans;
    }
}

C++

class Solution {
public:
    int numberOfArithmeticSlices(vector<int>& nums) {
        int ans = 0, cnt = 0;
        int d = 3000;
        for (int i = 0; i < nums.size() - 1; ++i) {
            if (nums[i + 1] - nums[i] == d) {
                ++cnt;
            } else {
                d = nums[i + 1] - nums[i];
                cnt = 0;
            }
            ans += cnt;
        }
        return ans;
    }
};

Go

func numberOfArithmeticSlices(nums []int) (ans int) {
	cnt, d := 0, 3000
	for i, b := range nums[1:] {
		a := nums[i]
		if b-a == d {
			cnt++
		} else {
			d = b - a
			cnt = 0
		}
		ans += cnt
	}
	return
}

TypeScript

function numberOfArithmeticSlices(nums: number[]): number {
    let ans = 0;
    let cnt = 0;
    let d = 3000;
    for (let i = 0; i < nums.length - 1; ++i) {
        const a = nums[i];
        const b = nums[i + 1];
        if (b - a == d) {
            ++cnt;
        } else {
            d = b - a;
            cnt = 0;
        }
        ans += cnt;
    }
    return ans;
}