Description
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.
Example 1:
Input: nums = [0,1,4,6,7,10], diff = 3
Output: 2
Explanation:
(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.
Example 2:
Input: nums = [4,5,6,7,8,9], diff = 2
Output: 2
Explanation:
(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.
Constraints:
3 <= nums.length <= 200
0 <= nums[i] <= 200
1 <= diff <= 50
nums is strictly increasing.
Solutions
Solution 1: Brute Force
We notice that the length of the array nums is no more than 200. Therefore, we can directly enumerate i, j, k, and check whether they meet the conditions. If they do, we increment the count of the triplet.
The time complexity is O(n3), where n is the length of the array nums. The space complexity is O(1).
PythonJavaC++GoTypeScript
class Solution:
def arithmeticTriplets(self, nums: List[int], diff: int) -> int:
return sum(b - a == diff and c - b == diff for a, b, c in combinations(nums, 3))(code-box)
class Solution {
public int arithmeticTriplets(int[] nums, int diff) {
int ans = 0;
int n = nums.length;
for (int i = 0; i < n; ++i) {
for (int j = i + 1; j < n; ++j) {
for (int k = j + 1; k < n; ++k) {
if (nums[j] - nums[i] == diff && nums[k] - nums[j] == diff) {
++ans;
}
}
}
}
return ans;
}
}(code-box)
class Solution {
public:
int arithmeticTriplets(vector<int>& nums, int diff) {
int ans = 0;
int n = nums.size();
for (int i = 0; i < n; ++i) {
for (int j = i + 1; j < n; ++j) {
for (int k = j + 1; k < n; ++k) {
if (nums[j] - nums[i] == diff && nums[k] - nums[j] == diff) {
++ans;
}
}
}
}
return ans;
}
};(code-box)
func arithmeticTriplets(nums []int, diff int) (ans int) {
n := len(nums)
for i := 0; i < n; i++ {
for j := i + 1; j < n; j++ {
for k := j + 1; k < n; k++ {
if nums[j]-nums[i] == diff && nums[k]-nums[j] == diff {
ans++
}
}
}
}
return
}(code-box)
function arithmeticTriplets(nums: number[], diff: number): number {
const n = nums.length;
let ans = 0;
for (let i = 0; i < n; ++i) {
for (let j = i + 1; j < n; ++j) {
for (let k = j + 1; k < n; ++k) {
if (nums[j] - nums[i] === diff && nums[k] - nums[j] === diff) {
++ans;
}
}
}
}
return ans;
}(code-box)
Solution 2: Array or Hash Table
We can first store the elements of nums in a hash table or array vis. Then, for each element x in nums, we check if x+diff and x+diff+diff are also in vis. If they are, we increment the count of the triplet.
After the enumeration, we return the answer.
The time complexity is O(n) and the space complexity is O(n), where n is the length of the array nums.
PythonJavaC++GoTypeScript
class Solution:
def arithmeticTriplets(self, nums: List[int], diff: int) -> int:
vis = set(nums)
return sum(x + diff in vis and x + diff * 2 in vis for x in nums)(code-box)
class Solution {
public int arithmeticTriplets(int[] nums, int diff) {
boolean[] vis = new boolean[301];
for (int x : nums) {
vis[x] = true;
}
int ans = 0;
for (int x : nums) {
if (vis[x + diff] && vis[x + diff + diff]) {
++ans;
}
}
return ans;
}
}(code-box)
class Solution {
public:
int arithmeticTriplets(vector<int>& nums, int diff) {
bitset<301> vis;
for (int x : nums) {
vis[x] = 1;
}
int ans = 0;
for (int x : nums) {
ans += vis[x + diff] && vis[x + diff + diff];
}
return ans;
}
};(code-box)
func arithmeticTriplets(nums []int, diff int) (ans int) {
vis := [301]bool{}
for _, x := range nums {
vis[x] = true
}
for _, x := range nums {
if vis[x+diff] && vis[x+diff+diff] {
ans++
}
}
return
}(code-box)
function arithmeticTriplets(nums: number[], diff: number): number {
const vis: boolean[] = new Array(301).fill(false);
for (const x of nums) {
vis[x] = true;
}
let ans = 0;
for (const x of nums) {
if (vis[x + diff] && vis[x + diff + diff]) {
++ans;
}
}
return ans;
}(code-box)