# [413] 等差数列划分

如果一个数列 至少有三个元素 ,并且任意两个相邻元素之差相同,则称该数列为等差数列。

例如,[1,3,5,7,9]、[7,7,7,7] 和 [3,-1,-5,-9] 都是等差数列。

给你一个整数数组 nums ,返回数组 nums 中所有为等差数组的 子数组 个数。

子数组 是数组中的一个连续序列。

示例 1:

输入:nums = [1,2,3,4]

输出:3

解释:nums 中有三个子等差数组:[1, 2, 3]、[2, 3, 4] 和 [1,2,3,4] 自身。

示例 2:

输入:nums = [1]

输出:0

这道题如果用动态规划的思路来做,可以设 dp[i][j] 表示从 i 到 j 的子数列是否为等差数列,若是的话存储等差数列的差值。之后我们来求解。

通过对题目分析我们注意到:

  1. j - i === 2 时,当 nums[j] - nums[j-1] === nums[j-1] - nums[j-2] 时,即表示这三个数组成的数列为等差数列,此时更新差值。dp[i][j] = nums[j] - nums[j-1]。
  2. 其余情况下,当 nums[j] - nums[j-1] === dp[i][j-1] 时,即新加入的数字差与之前的数列等差时,当前数列也为等差数列,更新差值。dp[i][j] = nums[j] - nums[j-1]。

而这就是我们的基准值与状态转移方程。

function numberOfArithmeticSlices(nums: number[]): number {
  let res = 0;
  const len = nums.length;
  const dp: (number | null)[][] = Array(len)
    .fill(0)
    .map(x => Array(len).fill(null));

  for (let i = 0; i < len; i++) {
    for (let j = i + 2; j < len; j++) {
      if (j - i === 2) {
        if (nums[j] - nums[j - 1] === nums[j - 1] - nums[j - 2]) {
          dp[i][j] = nums[j] - nums[j - 1];
          res += 1;
        }
      } else {
        if (dp[i][j - 1] !== null && nums[j] - nums[j - 1] === dp[i][j - 1]) {
          dp[i][j] = nums[j] - nums[j - 1];
          res += 1;
        }
      }
    }
  }
  return res;
}

而这样的空间复杂度明显太高了,会爆栈。因此我们需要进行简化。

我们可以注意到,事实上 dp[i][j] 的值并不依赖任何其他子问题,因此完全可以用一个临时变量 diffValue 来作为差值的存储。此时的空间复杂度仅为 O(1)。

接下来仅需判断新加入的元素差值是否与 diffValue 相同即可。

另外,我们还可以对遍历进行剪枝。当发现从 i 到 j 的子数列已经不是等差数列后,后面的数列自然也不会是等差数列,可以直接剪枝。

function numberOfArithmeticSlices(nums: number[]): number {
  let res = 0;
  const len = nums.length;

  for (let i = 0; i < len; i++) {
    let diffValue: number | undefined;
    for (let j = i + 2; j < len; j++) {
      if (j - i === 2) {
        if (nums[j] - nums[j - 1] === nums[j - 1] - nums[j - 2]) {
          diffValue = nums[j] - nums[j - 1];
          res += 1;
        }
      } else {
        if (diffValue === undefined || nums[j] - nums[j - 1] !== diffValue) {
          // 当 i 到 j 不为等差数列后,后面的数字也不用继续判断了
          break;
        } else {
          res += 1;
        }
      }
    }
  }
  return res;
}