LeetCode 413: Arithmetic Slices

Problem Restatement

We are given an integer array nums.

An arithmetic array is a subarray with at least 3 elements where the difference between consecutive elements is the same.

We must return the total number of arithmetic subarrays inside nums.

A subarray must be contiguous.

The official examples include [1,2,3,4] -> 3 and [1] -> 0. The constraints include 1 <= nums.length <= 5000 and -1000 <= nums[i] <= 1000. (github.com)

Input and Output

Example function shape:

def numberOfArithmeticSlices(nums: list[int]) -> int:
    ...

Examples

Example 1:

nums = [1, 2, 3, 4]

Arithmetic subarrays:

[1,2,3]
[2,3,4]
[1,2,3,4]

The answer is:

3

Example 2:

nums = [1]

The answer is:

0

because an arithmetic slice must contain at least 3 elements.

Example 3:

nums = [1, 3, 5, 7, 9]

Every consecutive difference is:

2

Arithmetic subarrays:

[1,3,5]
[3,5,7]
[5,7,9]
[1,3,5,7]
[3,5,7,9]
[1,3,5,7,9]

The answer is:

6

First Thought: Check Every Subarray

One direct approach is:

1. Generate every subarray of length at least 3.
2. Check whether consecutive differences are equal.
3. Count valid subarrays.

This works, but it is slow.

If the array length is n, there are:

O(n^2)

subarrays, and checking each one may take additional time.

We need a more efficient way.

Key Insight

Suppose we already know:

nums[i-2], nums[i-1], nums[i]

form an arithmetic sequence.

Then any arithmetic slice ending at:

i - 1

can be extended by:

nums[i]

if the difference stays the same.

Example:

1, 2, 3

is arithmetic.

When we append 4, we get:

1, 2, 3, 4

which is also arithmetic.

Additionally:

2, 3, 4

is a new arithmetic slice.

So every time the arithmetic condition continues, the number of new slices increases.

Dynamic Programming Idea

Define:

dp[i]

as:

The number of arithmetic slices ending at index i.

If:

nums[i] - nums[i-1] == nums[i-1] - nums[i-2]

then:

dp[i] = dp[i-1] + 1

Why?

Because:

1. Every arithmetic slice ending at i - 1 can extend to i.
2. The new length-3 slice ending at i also counts.

If the differences are not equal:

dp[i] = 0

The final answer is:

sum(dp)

Algorithm

Initialize:

total = 0
current = 0

Loop from index 2 to the end.

For each index i:

1. Compare consecutive differences.
2. If equal:
   • Increment current.
   • Add current to total.
3. Otherwise:
   • Reset current = 0.

Return total.

We can optimize the DP array into one variable because we only need the previous value.

Correctness

Suppose the consecutive differences at index i are equal:

nums[i] - nums[i-1]
=
nums[i-1] - nums[i-2]

Then:

nums[i-2], nums[i-1], nums[i]

forms a new arithmetic slice of length 3.

Also, every arithmetic slice ending at i - 1 can extend by nums[i] while preserving equal differences.

Therefore, if there were current arithmetic slices ending at i - 1, then there are:

current + 1

arithmetic slices ending at i.

The algorithm updates current exactly this way.

If the differences are unequal, no arithmetic slice can continue through index i, so resetting current = 0 is correct.

The algorithm adds every arithmetic slice exactly once, at the position where it ends.

Therefore, the final total is correct.

Complexity

Here n is the length of nums.

Implementation

from typing import List

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        total = 0
        current = 0

        for i in range(2, len(nums)):
            if nums[i] - nums[i - 1] == nums[i - 1] - nums[i - 2]:
                current += 1
                total += current
            else:
                current = 0

        return total

Code Explanation

We store:

current

which means:

The number of arithmetic slices ending at the previous index.

Initially:

current = 0

because arrays with fewer than 3 elements cannot form arithmetic slices.

Then we scan from index 2 onward:

for i in range(2, len(nums)):

We compare the two consecutive differences:

nums[i] - nums[i - 1]
==
nums[i - 1] - nums[i - 2]

If they are equal, arithmetic continuity continues.

So:

current += 1

The +1 counts the new length-3 slice.

Longer slices are extended automatically.

Then:

total += current

adds all arithmetic slices ending at index i.

If the differences are not equal:

current = 0

because no arithmetic slice can continue through this index.

Finally:

return total

returns the total number of arithmetic slices.

Testing

def test_arithmetic_slices():
    s = Solution()

    assert s.numberOfArithmeticSlices([1, 2, 3, 4]) == 3

    assert s.numberOfArithmeticSlices([1]) == 0

    assert s.numberOfArithmeticSlices([1, 3, 5, 7, 9]) == 6

    assert s.numberOfArithmeticSlices([7, 7, 7, 7]) == 3

    assert s.numberOfArithmeticSlices([1, 2, 4, 6, 8]) == 3

    assert s.numberOfArithmeticSlices([1, 1, 2, 5, 7]) == 0

    assert s.numberOfArithmeticSlices([3, -1, -5, -9]) == 3

    print("all tests passed")

Test Notes