2367. 等差三元组的数目#

题目描述#

给你一个下标从 0 开始、严格递增的整数数组 nums 和一个正整数 diff。如果满足下述全部条件，则三元组 (i, j, k) 就是一个等差三元组：

i < j < k，
nums[j] - nums[i] == diff，
nums[k] - nums[j] == diff。

返回数组 nums 中等差三元组的数目。

示例 1：

1
输入：nums = [0,1,4,6,7,10], diff = 3
2
输出：2
3
解释：(1, 2, 4) 是等差三元组：7 - 4 == 3 且 4 - 1 == 3。
4
     (2, 4, 5) 是等差三元组：10 - 7 == 3 且 7 - 4 == 3。

示例 2：

1
输入：nums = [4,5,6,7,8,9], diff = 2
2
输出：2
3
解释：(0, 2, 4) 是等差三元组：8 - 6 == 2 且 6 - 4 == 2。
4
     (1, 3, 5) 是等差三元组：9 - 7 == 2 且 7 - 5 == 2。

提示：

3 <= nums.length <= 200
0 <= nums[i] <= 200
1 <= diff <= 50
nums 严格递增

解题思路#

核心思想#

数组严格递增，因此三个指针 i < j < k 对应的值单调上升，且一旦 nums[j] - nums[i] > diff，随着 i 增大差值单调不增。这天然适合用三指针线性扫描。

枚举三元组最大的 k：用指针 j 找到「与 nums[k] 的差为 diff」的中间元素（即 nums[k] - nums[j] == diff），再用指针 i 找到「与 nums[j] 的差为 diff」的最小元素（nums[j] - nums[i] == diff）。由于三个值各自相差 diff，找到一次就构成一个等差三元组。

三指针同步推进示意（nums=[0,1,4,6,7,10], diff=3）：

1
k=4(nums[4]=7)时:
2
  让 j 前进，直到 nums[4]-nums[j] <= 3
3
    nums[4]-nums[2]=7-4=3 → j=2，差==diff
4
  再让 i 前进，直到 nums[2]-nums[i] <= 3
5
    nums[2]-nums[1]=4-1=3 → i=1，差==diff
6
  计数 +1  → 三元组 (1,2,4)
7
k=5(nums[5]=10)时:
8
  让 j 前进，直到 nums[5]-nums[j] <= 3 → j=4(nums[4]=7), 差=3
9
  让 i 前进，直到 nums[4]-nums[i] <= 3 → i=2(nums[2]=4), 差=3
10
  计数 +1  → 三元组 (2,4,5)
11
总数 = 2

算法步骤#

初始化指针 i = 0, j = 0，答案 ans = 0。
遍历 k 从 0 到 n-1，记 x = nums[k]：
- 内层一：while j < k and x - nums[j] > diff 时 j++；
- 若 j == k 或 x - nums[j] != diff，说明没有「与 x 相差 diff」的 j，跳过；
- 内层二：while i < j and nums[j] - nums[i] > diff 时 i++；
- 若 i < j 且 nums[j] - nums[i] == diff，说明发现 (i, j, k) 等差三元组，ans++。
返回 ans。

逐步举例#

以 nums=[0,1,4,6,7,10], diff=3 为例，只展示需计数的两步：

k	x=nums[k]	找 j（x-nums[j]==3）	找 i（nums[j]-nums[i]==3）	计数
4	7	nums[2]=4 → 差3，j=2	nums[1]=1 → 差3，i=1	+1
5	10	nums[4]=7 → 差3，j=4	nums[2]=4 → 差3，i=2	+1

中间值（4、7）未被重复计入，各被扫描常数次，源于严格递增数组保证的单调性。

复杂度分析#

时间复杂度：O(N)，i、j、k 三个指针各自单调右移，总体线性。
空间复杂度：O(1)，仅使用常数额外空间。

代码实现#

1
from typing import List
2

3

4
class Solution:
5
    def arithmeticTriplets(self, nums: List[int], diff: int) -> int:
6
        n = len(nums)
7
        i = j = 0
8
        ans = 0
9
        for k in range(n):
10
            x = nums[k]
11
            while j < k and x - nums[j] > diff:
12
                j += 1
13
            if j == k or x - nums[j] != diff:
14
                continue
15
            while i < j and nums[j] - nums[i] > diff:
16
                i += 1
17
            if i < j and nums[j] - nums[i] == diff:
18
                ans += 1
19
        return ans

1
import sys
2
from pathlib import Path
3

4
sys.path.insert(0, str(Path(__file__).parent))
5
import unittest
6

7
from solution import Solution
8

9

10
class TestSolution(unittest.TestCase):
11
    def test_example1(self):
12
        self.assertEqual(Solution().arithmeticTriplets([0, 1, 4, 6, 7, 10], 3), 2)
13

14
    def test_example2(self):
15
        self.assertEqual(Solution().arithmeticTriplets([4, 5, 6, 7, 8, 9], 2), 2)
16

17

18
if __name__ == '__main__':
19
    unittest.main()