You are given a 0-indexed, strictly increasing integer array nums
and a positive integer diff
. A triplet (i, j, k)
is an arithmetic triplet if the following conditions are met:
i < j < k
,nums[j] - nums[i] == diff
, andnums[k] - nums[j] == diff
.Return the number of unique arithmetic triplets.
For example:
nums = [0,1,4,6,7,10], diff = 3
Output: 2. (1, 2, 4) is an arithmetic triplet because both 7 - 4 == 3 and 4 - 1 == 3. (2, 4, 5) is an arithmetic triplet because both 10 - 7 == 3 and 7 - 4 == 3.
nums = [4,5,6,7,8,9], diff = 2
Output: 2. (0, 2, 4) is an arithmetic triplet because both 8 - 6 == 2 and 6 - 4 == 2. (1, 3, 5) is an arithmetic triplet because both 9 - 7 == 2 and 7 - 5 == 2.
## Naive Solution
The most straightforward approach is to iterate through all possible triplets (i, j, k) and check if they satisfy the arithmetic triplet conditions.
```python
def count_arithmetic_triplets_naive(nums, diff):
count = 0
n = len(nums)
for i in range(n):
for j in range(i + 1, n):
for k in range(j + 1, n):
if nums[j] - nums[i] == diff and nums[k] - nums[j] == diff:
count += 1
return count
We can optimize this by using a set to store the numbers we have seen so far. Then, for each number nums[j]
, we check if nums[j] - diff
and nums[j] + diff
are present in the set. This significantly reduces the time complexity.
def count_arithmetic_triplets_optimal(nums, diff):
count = 0
seen = set()
for num in nums:
if (num - diff) in seen and (num + diff) in seen:
count += 1
seen.add(num)
return count
The naive solution has three nested loops, each iterating up to n
times, where n
is the length of the nums
array. Therefore, the time complexity is O(n^3).
The optimal solution iterates through the nums
array once (O(n)). Inside the loop, we perform constant-time set lookups (O(1)). Therefore, the overall time complexity is O(n).
The naive solution does not use any extra space that scales with the input size. It uses a constant amount of extra space for variables. Therefore, the space complexity is O(1).
The optimal solution uses a set to store the numbers we have seen so far. In the worst case, we might store all the numbers in the nums
array in the set. Therefore, the space complexity is O(n).
nums
is empty or has less than 3 elements, there can be no arithmetic triplets. The code handles this implicitly, as the loops won't execute.nums
is strictly increasing, so there won't be duplicate numbers.diff
is very large and the numbers in nums
are relatively small, it's possible that no arithmetic triplets exist. The code will correctly return 0 in this case.0 <= nums[i] <= 200
and 1 <= diff <= 50
, there is no risk of integer overflow when calculating nums[j] - nums[i]
or nums[k] - nums[j]
.