# [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 的子数列是否为等差数列,若是的话存储等差数列的差值。之后我们来求解。
通过对题目分析我们注意到:
- j - i === 2 时,当 nums[j] - nums[j-1] === nums[j-1] - nums[j-2] 时,即表示这三个数组成的数列为等差数列,此时更新差值。dp[i][j] = nums[j] - nums[j-1]。
- 其余情况下,当 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;
}