0413-等差数列划分

Raphael Liu Lv10
  2023-09-13 00:47:26 2023-09-13 00:47:26 Created   2023-09-13 15:32:49 2023-09-13 15:32:49 Updated    

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

  • 例如，[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

提示：

  • 1 <= nums.length <= 5000
  • -1000 <= nums[i] <= 1000

方法一：差分 + 计数

思路与算法

考虑一个比较直观的做法：

  • 我们枚举等差数列的最后两项 nums}[i - 1]$ 以及 nums}[i]$，那么等差数列的公差 $d$ 即为 nums}[i - 1] - \textit{nums}[i]$；

  • 随后我们使用一个指针 $j$ 从 $i - 2$ 开始逆序地遍历数组的前缀部分 nums}[0 .. i-2]$：

    • 如果 nums}[j] - \textit{nums}[j + 1] = d$，那么说明 nums}[j], \cdots, \textit{nums}[i]$ 组成了一个长度至少为 $3$ 的等差数列，答案增加 $1$；

    • 否则更小的 $j$ 也无法作为等差数列的首个位置了，我们直接退出遍历。

这个做法的时间复杂度是 $O(n^2)$ 的，即枚举最后两项的时间复杂度为 $O(n)$，使用指针 $j$ 遍历的时间复杂度也为 $O(n)$，相乘得到总时间复杂度 $O(n^2)$。对于一些运行较慢的语言，该方法可能会超出时间限制，因此我们需要进行优化。

优化

如果我们已经求出了 nums}[i - 1]$ 以及 nums}[i]$ 作为等差数列的最后两项时，答案增加的次数 $t_i$，那么能否快速地求出 $t_{i+1 呢？

答案是可以的：

  • 如果 nums}[i] - \textit{nums}[i + 1] = d$，那么在这一轮遍历中，$j$ 会遍历到与上一轮相同的位置，答案增加的次数相同，并且额外多出了 nums}[i-1], \textit{nums}[i], \textit{nums}[i+1]$ 这一个等差数列，因此有：

$$
t_{i+1} = t_i + 1
$$

  • 如果 nums}[i] - \textit{num}[i + 1] \neq d$，那么 $j$ 从初始值 $i-1$ 开始就会直接退出遍历，答案不会增加，因此有：

$$
t_{i+1} = 0
$$

这样一来，我们通过上述简单的递推，即可在 $O(1)$ 的时间计算等差数列数量的增量，总时间复杂度减少至 $O(n)$。

代码

[sol1-C++]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
public:
    int numberOfArithmeticSlices(vector<int>& nums) {
        int n = nums.size();
        if (n == 1) {
            return 0;
        }

        int d = nums[0] - nums[1], t = 0;
        int ans = 0;
        // 因为等差数列的长度至少为 3，所以可以从 i=2 开始枚举
        for (int i = 2; i < n; ++i) {
            if (nums[i - 1] - nums[i] == d) {
                ++t;
            }
            else {
                d = nums[i - 1] - nums[i];
                t = 0;
            }
            ans += t;
        }
        return ans;
    }
};
[sol1-Java]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution {
    public int numberOfArithmeticSlices(int[] nums) {
        int n = nums.length;
        if (n == 1) {
            return 0;
        }

        int d = nums[0] - nums[1], t = 0;
        int ans = 0;
        // 因为等差数列的长度至少为 3，所以可以从 i=2 开始枚举
        for (int i = 2; i < n; ++i) {
            if (nums[i - 1] - nums[i] == d) {
                ++t;
            } else {
                d = nums[i - 1] - nums[i];
                t = 0;
            }
            ans += t;
        }
        return ans;
    }
}
[sol1-C#]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
public class Solution {
    public int NumberOfArithmeticSlices(int[] nums) {
        int n = nums.Length;
        if (n == 1) {
            return 0;
        }

        int d = nums[0] - nums[1], t = 0;
        int ans = 0;
        // 因为等差数列的长度至少为 3，所以可以从 i=2 开始枚举
        for (int i = 2; i < n; ++i) {
            if (nums[i - 1] - nums[i] == d) {
                ++t;
            } else {
                d = nums[i - 1] - nums[i];
                t = 0;
            }
            ans += t;
        }
        return ans;
    }
}
[sol1-Python3]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        n = len(nums)
        if n == 1:
            return 0
        
        d, t = nums[0] - nums[1], 0
        ans = 0
        
        # 因为等差数列的长度至少为 3，所以可以从 i=2 开始枚举
        for i in range(2, n):
            if nums[i - 1] - nums[i] == d:
                t += 1
            else:
                d = nums[i - 1] - nums[i]
                t = 0
            ans += t
        
        return ans
[sol1-Golang]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
func numberOfArithmeticSlices(nums []int) (ans int) {
    n := len(nums)
    if n == 1 {
        return
    }

    d, t := nums[0]-nums[1], 0
    // 因为等差数列的长度至少为 3，所以可以从 i=2 开始枚举
    for i := 2; i < n; i++ {
        if nums[i-1]-nums[i] == d {
            t++
        } else {
            d, t = nums[i-1]-nums[i], 0
        }
        ans += t
    }
    return
}
[sol1-C]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
int numberOfArithmeticSlices(int* nums, int numsSize) {
    if (numsSize == 1) {
        return 0;
    }

    int d = nums[0] - nums[1], t = 0;
    int ans = 0;
    // 因为等差数列的长度至少为 3，所以可以从 i=2 开始枚举
    for (int i = 2; i < numsSize; ++i) {
        if (nums[i - 1] - nums[i] == d) {
            ++t;
        } else {
            d = nums[i - 1] - nums[i];
            t = 0;
        }
        ans += t;
    }
    return ans;
}
[sol1-JavaScript]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
var numberOfArithmeticSlices = function(nums) {
    const n = nums.length;
    if (n === 1) {
        return 0;
    }

    let d = nums[0] - nums[1], t = 0;
    let ans = 0;
    // 因为等差数列的长度至少为 3，所以可以从 i=2 开始枚举
    for (let i = 2; i < n; ++i) {
        if (nums[i - 1] - nums[i] === d) {
            ++t;
        } else {
            d = nums[i - 1] - nums[i];
            t = 0;
        }
        ans += t;
    }
    return ans;
};

复杂度分析

  • 时间复杂度：$O(n)$，其中 $n$ 是数组 nums 的长度。

  • 空间复杂度：$O(1)$。

 Comments
On this page
0413-等差数列划分