Problem statement

https://leetcode.com/problems/arithmetic-slices-ii-subsequence/

Solution

Let dp[i][d] denotes the number of arithmetic subsequences that ends with A[i] and its common difference is d. Then we can look at all j < i and check if we can continue subsequence ending with point j. Then difference is equal to d = A[i] - A[j] and we can update dp[i][d] += (dp[j][d] + 1). Notice, that in this way we evaluate number of all possible arithmetic subsequences with length >= 2, so we just need to subtract n(n-1)/2 from final number (be careful to evaluate sum for dp[i], after we iterate over all j).

Complexity

Time and space complexity is O(n^2) for double loop and to keep all dictionaries.

Code

class Solution:
    def numberOfArithmeticSlices(self, A):
        total, n = 0, len(A)
        dp = [Counter() for item in A]
        for i in range(n):
            for j in range(i):
                dp[i][A[i] - A[j]] += (dp[j][A[i] - A[j]] + 1)      
            total += sum(dp[i].values())
          
        return total - (n-1)*n//2