Taro Logo

Number of Arithmetic Triplets

Easy
2 months ago

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, and
  • nums[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.

Sample Answer
## 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

Optimal Solution

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

Big(O) Run-time Analysis

Naive Solution:

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).

Optimal Solution:

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).

Big(O) Space Usage Analysis

Naive Solution:

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).

Optimal Solution:

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).

Edge Cases

  1. Empty Array: If the input array 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.
  2. Duplicate Numbers: The problem statement specifies that nums is strictly increasing, so there won't be duplicate numbers.
  3. Large Diff: If 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.
  4. Integer Overflow: Since the constraints specify that 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].