Dynamic Programming Algorithm to Count Contiguous Arithmetic Sequences

Given a list of integers nums, return the number of contiguous arithmetic sequences of length ≥ 3.

Constraints
n ≤ 10,000 where n is length of nums.
Example 1
Input
nums = [5, 7, 9, 11, 12, 13]
Output
4
Explanation
We have the following arithmetic sequences:

[5, 7, 9]
[7, 9, 11]
[5, 7, 9, 11]
[11, 12, 13]

Bruteforce Algorithm to Count the Arithmetic Sequences

We can bruteforce – checking each possible sequence which takes O(N^2) quadratic time, and then validate each sequence that will take O(N) time. Overall, the time complexity is O(N^3).

int countArithmeticSequences(vector<int>& nums) {
    int n = static_cast<int>(nums.size());
    if (n == 0) return 0;
    int ans = 0;
    function<bool(int, int)> isGood = [&](int a, int b) {
        for (int i = a + 1; i <= b; ++ i) {
            if (nums[i] - nums[i - 1] != nums[a + 1] - nums[a]) {
                return false;
            }
        }
        return true;
    };
    for (int i = 0; i < n; ++ i) {
        for (int j = i + 2; j < n; ++ j) {
            // each subsequence [i, j]
            if (isGood(i, j)) { // check each subsequence
                ans ++;
            }
        }
    }
    return ans;
}

Count the Arithmetic Sequences via DP Algorithm

We can accumulate the counting using DP array. For example, dp[i] meaning the number of arithmetic sequences that end at index i. Then, we accumulate the dp[i] array when we find subsequence of at least 3 numbers that satisfy the arithmetic sequences.

C++ Dynamic Programming Algorithm to count the arithmetic sequences:

int countArithmeticSequences(vector<int>& nums) {
    int n = static_cast<int>(nums.size());
    if (n == 0) return 0;
    vector<int> dp(n, 0);
    int ans = 0;
    for (int i = 2; i < n; ++ i) {
        if (nums[i] - nums[i - 1] == nums[i - 1] - nums[i - 2]) {
            dp[i] = dp[i - 1] + 1;
            ans += dp[i];
        }
    }
    return ans;
}

The python implementation of the same algorithm:

class Solution:
    def countArithmeticSequences(self, nums):
        ans = 0
        n = len(nums)
        dp = [0] * n
        for i in range(2, n):
            if nums[i - 1] - nums[i - 2] == nums[i] - nums[i - 1]:
                dp[i] = dp[i - 1] + 1
                ans += dp[i]
        return ans

Both implementation have O(N) time complexity and O(N) space.

See also: Can We Make Arithmetic Progression From Sequence of Numbers?

–EOF (The Ultimate Computing & Technology Blog) —