Count Subarrays of Length Three With a Condition
EasyPrompt
Given an integer array nums, return the number of subarrays of length 3 such that the sum of the first and third numbers equals exactly half of the second number.
Example 1:
Input: nums = [1,2,1,4,1]
Output: 1
Explanation:
Only the subarray [1,4,1] contains exactly 3 elements where the sum of the first and third numbers equals half the middle number.
Example 2:
Input: nums = [1,1,1]
Output: 0
Explanation:
[1,1,1] is the only subarray of length 3. However, its first and third numbers do not add to half the middle number.
Constraints:
3 <= nums.length <= 100-100 <= nums[i] <= 100
Approaches
2 approaches with complexity analysis and trade-offs.
This approach iterates through all possible combinations of three indices (i, j, k) in the array and checks if they form a valid subarray of length 3 that satisfies the given condition.
Algorithm
- Initialize a counter
countto 0. - Iterate with an index
ifrom0tonums.length - 1. - Inside, iterate with an index
jfrom0tonums.length - 1. - Inside, iterate with an index
kfrom0tonums.length - 1. - If
j == i + 1andk == i + 2: - If
2 * (nums[i] + nums[k]) == nums[j]: - Increment
count. - Return
count.
Walkthrough
The core idea is to use three nested loops to pick three elements from the array at indices i, j, and k. Inside the innermost loop, we first validate if these indices actually form a contiguous subarray of length 3 by checking if j == i + 1 and k == i + 2. If they do, we then check if they satisfy the problem's condition. To avoid floating-point arithmetic, the condition nums[i] + nums[k] == nums[j] / 2 is rewritten as 2 * (nums[i] + nums[k]) == nums[j]. A counter is incremented for each valid subarray found.
class Solution { public int countSubarrays(int[] nums) { int n = nums.length; int count = 0; if (n < 3) { return 0; } for (int i = 0; i < n; i++) { for (int j = 0; j < n; j++) { for (int k = 0; k < n; k++) { // Check if the indices form a contiguous subarray of length 3 if (j == i + 1 && k == i + 2) { // Check the condition // Using 2 * (a + c) == b to avoid floating point issues if (2 * (nums[i] + nums[k]) == nums[j]) { count++; } } } } } return count; }}Complexity
Time
O(n^3), where n is the length of the `nums` array. This is because of the three nested loops, each iterating up to `n` times.
Space
O(1), as we only use a constant amount of extra space for the counter and loop variables.
Trade-offs
Pros
Conceptually simple to understand, as it directly translates the problem of finding three elements into three loops.
Cons
Highly inefficient due to the cubic time complexity. It performs a lot of unnecessary checks for indices that do not form a contiguous subarray.
Not practical for larger input sizes, although it passes for the given constraints (n <= 100).
A much more efficient approach is to iterate through the array just once, considering each possible starting position of a subarray of length 3. This avoids the redundant checks of the brute-force method.
Algorithm
- Initialize a counter
countto 0. - Iterate with an index
ifrom0tonums.length - 3. - Define the three elements of the subarray:
first = nums[i],second = nums[i+1],third = nums[i+2]. - Check if
2 * (first + third) == second. - If the condition is true, increment
count. - After the loop, return
count.
Walkthrough
This optimal approach recognizes that a subarray of length 3 is uniquely defined by its starting position. We can iterate through the array with a single loop, where the loop variable i represents the starting index of a subarray. The loop runs from i = 0 to n - 3, where n is the array length, covering all possible subarrays of length 3. For each i, we consider the subarray [nums[i], nums[i+1], nums[i+2]]. We then check if it meets the condition 2 * (nums[i] + nums[i+2]) == nums[i+1]. This check avoids floating-point division and is performed in constant time. If the condition is true, we increment a counter. This single pass is sufficient to count all such subarrays.
class Solution { public int countSubarrays(int[] nums) { int n = nums.length; int count = 0; // A subarray of length 3 needs at least 3 elements. if (n < 3) { return 0; } // Iterate through all possible starting positions of a subarray of length 3. // The last possible starting index is n - 3. for (int i = 0; i <= n - 3; i++) { // The subarray is [nums[i], nums[i+1], nums[i+2]] int first = nums[i]; int second = nums[i+1]; int third = nums[i+2]; // Check the condition: first + third == second / 2 // To avoid floating point math, we check: 2 * (first + third) == second if (2 * (first + third) == second) { count++; } } return count; }}Complexity
Time
O(n), where n is the length of the `nums` array. We iterate through the array once, from index 0 to n-3. Inside the loop, we perform a constant number of operations.
Space
O(1), as we only use a constant amount of extra space for the counter and loop variable.
Trade-offs
Pros
Optimal time complexity. It's the most efficient way to solve the problem as we must examine each potential subarray of length 3.
Simple and clean implementation.
Cons
No significant cons for this problem, as it's the optimal solution.
Solutions
Solution
class Solution {public int countSubarrays(int[] nums) { int ans = 0; for (int i = 1; i + 1 < nums.length; ++i) { if ((nums[i - 1] + nums[i + 1]) * 2 == nums[i]) { ++ans; } } return ans; }}class Solution {public: int countSubarrays(vector<int> &nums) { int ans = 0; for (int i = 1; i + 1 < nums.size(); ++i) { if ((nums[i - 1] + nums[i + 1]) * 2 == nums[i]) { ++ans; } } return ans; }};class Solution: def countSubarrays(self, nums: List[int]) -> int: return sum( (nums[i - 1] + nums[i + 1]) * 2 == nums[i] for i in range(1, len(nums) - 1))Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.