0413. 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

Analysis

This question can be solved using math. We can try to solve for the subarray with a size of 3 first, and it will only form 1 arithmetic slice. Now we think of the subarray with the size of 4. We can split the size-4 subarray into 2 size-three subarrays:

1,2,3,4 -> 1,2,3 + 2,3,4 + 1,2,3,4
size-four -> size-three + size-three + 1 = 1 + 1 + 1 = 3

Now we can find some relationship between the current size subarray and its previous size subarray.

1,2,3,4,5 -> 1,2,3,4 + 2,3,4,5 + 1,2,3,4,5
size-five -> size-four + size-four + 1 = 7 (incorrect, should be 6 instead)

However, the above result is incorrect, and it should be 6 instead. What have we done wrong here? We double-count the 2,3,4 because it's included in both 1,2,3,4 and 2,3,4,5 subarrays.

1~n -> 1~n-1 + 2~n + 1~n - 2~n-1 (for n > 3)
size-n -> size-(n-1) + size-(n-1) - size-(n-2) + 1

So right now we can already start implementing this using: dp[i] = 2 * dp[i - 1] - dp[i - 2] + 1, but I have found another simpler equation: dp[i] = dp[i - 1] + i - 2 and I am unable to prove the correctness. This equation is from Leetcode discussion.

• Time: $O(n)$
• Space: $O(n)$

Code

Using my DP equation

class Solution {
public:
    int numberOfArithmeticSlices(vector<int>& nums) {
        int res = 0, n = nums.size();        
        map<int, int> dp; // cnt -> slices
        dp[3] = 1;        
        int i = 0;
        while (i + 1 < n) {
            int j = i + 1, diff = nums[j] - nums[j - 1];            
            while (j < n && nums[j] - nums[j - 1] == diff) j ++;
            int cnt = j - i;
            i = j - 1;
            if (dp.count(cnt)) res += dp[cnt];
            else {
                for (int i = 4; i <= cnt; ++i) {
                    if (i - 2 >= 3) // there is only duplication when middle subarray is longer than 3
                        dp[i] = 2 * dp[i - 1] - dp[i - 2] + 1;
                    else
                        dp[i] = 2 * dp[i - 1] + 1;
                }
                res += dp[cnt];
            }                

        }
        return res;
    }
};

Using DP equation from discussion

class Solution {
public:
    int numberOfArithmeticSlices(vector<int>& nums) {
        int res = 0, n = nums.size();        
        map<int, int> dp; // cnt -> slices
        dp[3] = 1;        
        int i = 0;
        while (i + 1 < n) {
            int j = i + 1, diff = nums[j] - nums[j - 1];            
            while (j < n && nums[j] - nums[j - 1] == diff) j ++;
            int cnt = j - i;
            i = j - 1;
            if (dp.count(cnt)) res += dp[cnt];
            else {
                for (int i = 4; i <= cnt; ++i) {
                    dp[i] = dp[i - 1] + i - 2;
                }
                res += dp[cnt];
            }                

        }
        return res;
    }
};