413.等差数列划分.md

413.等差数列划分

/*
 * @lc app=leetcode.cn id=413 lang=typescript
 *
 * [413] 等差数列划分
 */

// @lc code=start
function numberOfArithmeticSlices(nums: number[]): number {}
// @lc code=end

解法 1: 动态规划

• 状态: 第 i 个数字结尾时的连续子数组的索引起点
• 递推公式:
  • 当 i,i-1,i-2 能组成等差数列时: dp[i]=min(dp[i-1],i-2)
  • 否则: dp[i]=i
• 边界:
  • dp=[0,1]

使用 res 记录每段等差数列的长度,然后遍历 res,将每个长度的子集数相加,即为题目的结果

求长度为 n 的等差数列的连续子集:

• 3 -> 1
• 4 -> 3 -> 1 + 2
• 5 -> 6 -> 1 + 2 + 3
• ...
• n -> 1+...+(n-2) -> (n-2)*(n-3)/2

function numberOfArithmeticSlices(nums: number[]): number {
  let dp = new Array(nums.length).fill(0).map((_, i) => i)
  let res = []
  for (let i = 2; i <= nums.length; i++) {
    if (nums[i] - nums[i - 1] === nums[i - 1] - nums[i - 2]) {
      dp[i] = Math.min(dp[i - 1], i - 2)
    } else {
      res.push(i - dp[i - 1] + 1)
    }
  }

  return res.reduce((sum, n) => sum + ((n - 2) * (n - 3)) / 2, 0)
}

解法 2: 动态规划

• 状态:
  • t: 以第 i 个数字结尾时的等差数列个数
  • dp[i]: 1~i 之内所有的等差数列子数组个数
• 递推公式:
  • 当 i,i-1,i-2 能组成等差数列时, t++
  • 否则 t=0
  • dp[i]=dp[i-1]+t
• 边界:
  • dp=[0,0]
  • t=0

function numberOfArithmeticSlices(nums: number[]): number {
  let [dp, t] = [[0, 0], 0]
  for (let i = 2; i < nums.length; i++) {
    if (nums[i] - nums[i - 1] === nums[i - 1] - nums[i - 2]) t++
    else t = 0

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

  return dp[dp.length - 1]
}

优化空间

function numberOfArithmeticSlices(nums: number[]): number {
  let [count, t] = [0, 0]
  for (let i = 2; i < nums.length; i++) {
    if (nums[i] + nums[i - 2] === 2 * nums[i - 1]) t++
    else t = 0

    count += t
  }

  return count
}

Case

test.each([
  { input: { nums: [1, 2, 3, 4] }, output: 3 },
  { input: { nums: [1, 2, 3, 4, 5] }, output: 6 },
  { input: { nums: [1, 2, 3, 4, 5, 6] }, output: 10 },
  { input: { nums: [9, 2, 3, 4, 5, 10, 6, 7, 8] }, output: 4 },
  { input: { nums: [1] }, output: 0 },
])('input: nums = $input.nums', ({ input: { nums }, output }) => {
  expect(numberOfArithmeticSlices(nums)).toBe(output)
})