413. 等差数列划分(Medium)

题目描述

返回所有等差子数列的个数

样例

Input: A = [1, 2, 3, 4]
output: 3 
# 解释，等差子数列的个数 [1,2,3],[2,3,4],[1,2,3,4] 

题解

图文教程参考：LC官方

动态规划问题

转移方程

• 当 nums[i] 可以和 nums[i - 1]，nums[i - 2] 组成等差数列

  转移方程为 $dp[i] = dp[i - 1] + 1$

• 当 nums[i] 可以和 nums[i - 1]，nums[i - 2] 无法组成等差数列

  转移方程为 $dp[i] = 0$

举例说明，[1, 2, 3, 4, 5]

dp[2] = 1 ，对应的等差子序列为：[1, 2, 3]

dp[3] = dp[2] + 1=2，对应的等差子序列为：[1, 2, 3, 4]， [2, 3, 4] 。[2, 3, 4] 是加上的那个 1

dp[4] = dp[3] + 1=3，对应的等差子序列为：[1, 2, 3, 4, 5]， [2, 3, 4, 5], [3, 4, 5] 。[3, 4, 5] 是加上的那个 1

dp[i] 代表的是以 num[i] 为结尾的等差子序列的个数

代码

func sum(nums []int) int {
    ans := 0
    for _, num := range nums {
        ans += num 
    }
    return ans 
}

func numberOfArithmeticSlices(A []int) int {
    n := len(A)
    if n < 3 { return 0 }
    dp := make([]int, n)
    for i := 2; i < n; i++ {
        if A[i] - A[i - 1] == A[i - 1] - A[i - 2] {
            dp[i] = dp[i - 1] + 1
        } 
      	// 默认就是 0 
    }
    return sum(dp)
}