3392. Count Subarrays of Length Three With a Condition

Problem Description

Given an integer array nums, you need to find and count all subarrays of length 3 where a specific mathematical relationship exists between the three elements.

For each subarray of exactly 3 consecutive elements, check if the sum of the first element and the third element equals exactly half of the middle (second) element. In other words, for a subarray [a, b, c], the condition to satisfy is: a + c = b / 2.

The task is to return the total count of all such subarrays in the given array that meet this criterion.

For example, if we have a subarray [1, 4, 3], we would check if 1 + 3 = 4 / 2, which gives us 4 = 2. Since this is not true, this subarray would not be counted. However, if we had [1, 6, 2], we would check if 1 + 2 = 6 / 2, which gives us 3 = 3. Since this is true, this subarray would be counted.

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

Intuition

The problem asks us to examine all possible subarrays of length 3, which immediately suggests we need to look at every consecutive triplet in the array. Since we're dealing with fixed-size subarrays (always length 3), we can use a sliding window approach where we examine each position that could be the middle element of such a subarray.

The key insight is recognizing that for any valid middle position i, we need to check the elements at positions i-1 (first element), i (middle element), and i+1 (third element). The middle element can be any index from 1 to len(nums) - 2 because we need one element before it and one element after it to form a complete subarray of length 3.

The mathematical condition a + c = b / 2 can be rearranged to avoid floating-point division. Instead of checking if nums[i-1] + nums[i+1] = nums[i] / 2, we can multiply both sides by 2 to get (nums[i-1] + nums[i+1]) * 2 = nums[i]. This avoids potential precision issues with division and ensures we're working with integers throughout.

By iterating through all valid middle positions and checking this condition for each triplet, we can count how many subarrays satisfy the requirement. The single-pass nature of this approach makes it efficient with O(n) time complexity, as we only need to visit each potential middle element once.

Solution Approach

The implementation uses a single pass through the array to count all valid subarrays of length 3. Here's how the solution works:

Iterate through valid middle positions: We traverse from index 1 to len(nums) - 2. These are all the positions where an element can serve as the middle element of a subarray of length 3. Position 0 cannot be a middle element (no element before it), and position len(nums) - 1 cannot be a middle element (no element after it).
Check the condition for each triplet: For each middle position i, we examine the triplet consisting of:
- First element: nums[i - 1]
- Middle element: nums[i]
- Third element: nums[i + 1]
Verify the mathematical relationship: We check if (nums[i - 1] + nums[i + 1]) * 2 == nums[i]. This is equivalent to checking if the sum of the first and third elements equals half of the middle element, but avoids division by multiplying both sides by 2.
Count valid subarrays: The solution uses a generator expression with the sum() function. Each comparison (nums[i - 1] + nums[i + 1]) * 2 == nums[i] evaluates to True (counted as 1) or False (counted as 0). The sum() function adds up all the True values, giving us the total count of valid subarrays.

The algorithm processes each element at most once as a middle element, resulting in O(n) time complexity where n is the length of the input array. The space complexity is O(1) as we only use a constant amount of extra space for the iteration variable and temporary calculations.

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through the solution with the array nums = [2, 8, 3, 1, 4, 2].

We need to find all subarrays of length 3 where the sum of the first and third elements equals half of the middle element.

Step 1: Identify valid middle positions

Array length is 6, so valid middle positions are indices 1 to 4 (inclusive)
These positions allow for one element before and one element after

Step 2: Check each triplet

Position i = 1 (middle element = 8):

Triplet: [2, 8, 3] (indices 0, 1, 2)
Check: Does 2 + 3 = 8/2?
2 + 3 = 5 and 8/2 = 4
5 ≠ 4, so this doesn't count

Position i = 2 (middle element = 3):

Triplet: [8, 3, 1] (indices 1, 2, 3)
Check: Does 8 + 1 = 3/2?
8 + 1 = 9 and 3/2 = 1.5
9 ≠ 1.5, so this doesn't count

Position i = 3 (middle element = 1):

Triplet: [3, 1, 4] (indices 2, 3, 4)
Check: Does 3 + 4 = 1/2?
3 + 4 = 7 and 1/2 = 0.5
7 ≠ 0.5, so this doesn't count

Position i = 4 (middle element = 4):

Triplet: [1, 4, 2] (indices 3, 4, 5)
Check: Does 1 + 2 = 4/2?
1 + 2 = 3 and 4/2 = 2
3 ≠ 2, so this doesn't count

Step 3: Count total valid subarrays Total count = 0 (no subarrays satisfied the condition)

Now let's try another example where we do find valid subarrays: nums = [1, 6, 2, 10, 5]

Position i = 1: [1, 6, 2] → 1 + 2 = 3 and 6/2 = 3 → Valid! ✓ Position i = 2: [6, 2, 10] → 6 + 10 = 16 and 2/2 = 1 → Invalid Position i = 3: [2, 10, 5] → 2 + 5 = 7 and 10/2 = 5 → Invalid

Total count = 1

The implementation avoids division by rearranging the equation to (nums[i-1] + nums[i+1]) * 2 == nums[i], which for our valid case becomes (1 + 2) * 2 = 6 → 3 * 2 = 6 → 6 = 6 ✓

Solution Implementation

1class Solution:
2    def countSubarrays(self, nums: List[int]) -> int:
3        """
4        Count the number of subarrays of length 3 where the middle element
5        equals twice the sum of its neighbors.
6      
7        A subarray [nums[i-1], nums[i], nums[i+1]] is valid if:
8        nums[i] = 2 * (nums[i-1] + nums[i+1])
9      
10        Args:
11            nums: List of integers
12          
13        Returns:
14            Count of valid subarrays of length 3
15        """
16        # Initialize counter for valid subarrays
17        count = 0
18      
19        # Iterate through all possible middle positions
20        # Start from index 1 and end at len(nums) - 1 to ensure valid triplets
21        for i in range(1, len(nums) - 1):
22            # Check if middle element equals twice the sum of its neighbors
23            if nums[i] == 2 * (nums[i - 1] + nums[i + 1]):
24                count += 1
25      
26        return count
27

1class Solution {
2    /**
3     * Counts the number of subarrays of length 3 where the middle element
4     * equals twice the sum of its adjacent elements.
5     * 
6     * @param nums the input array of integers
7     * @return the count of valid subarrays
8     */
9    public int countSubarrays(int[] nums) {
10        int count = 0;
11      
12        // Iterate through all possible middle elements of subarrays of length 3
13        // Start from index 1 and end at index length-2 to ensure valid triplets
14        for (int middleIndex = 1; middleIndex < nums.length - 1; middleIndex++) {
15            // Check if the middle element equals twice the sum of its neighbors
16            // Condition: nums[middle] = 2 * (nums[left] + nums[right])
17            int leftElement = nums[middleIndex - 1];
18            int middleElement = nums[middleIndex];
19            int rightElement = nums[middleIndex + 1];
20          
21            if (middleElement == 2 * (leftElement + rightElement)) {
22                count++;
23            }
24        }
25      
26        return count;
27    }
28}
29

1class Solution {
2public:
3    int countSubarrays(vector<int>& nums) {
4        int count = 0;
5      
6        // Iterate through all possible middle elements of subarrays of length 3
7        // Start from index 1 and end at index n-2 to ensure valid triplets
8        for (int i = 1; i + 1 < nums.size(); ++i) {
9            // Check if the middle element equals half the sum of its neighbors
10            // This is equivalent to checking if nums[i] = (nums[i-1] + nums[i+1]) / 2
11            // Multiply by 2 to avoid floating point division
12            if (nums[i] * 2 == nums[i - 1] + nums[i + 1]) {
13                ++count;
14            }
15        }
16      
17        return count;
18    }
19};
20

1/**
2 * Counts the number of subarrays of length 3 where the middle element
3 * equals twice the average of its neighboring elements.
4 * 
5 * @param nums - The input array of numbers
6 * @returns The count of valid subarrays
7 */
8function countSubarrays(nums: number[]): number {
9    let count: number = 0;
10  
11    // Iterate through potential middle elements (index 1 to length-2)
12    for (let middleIndex: number = 1; middleIndex + 1 < nums.length; middleIndex++) {
13        const leftElement: number = nums[middleIndex - 1];
14        const middleElement: number = nums[middleIndex];
15        const rightElement: number = nums[middleIndex + 1];
16      
17        // Check if middle element equals twice the average of left and right elements
18        // Condition: (left + right) * 2 === middle is equivalent to:
19        // middle === 2 * (left + right) / 2
20        if ((leftElement + rightElement) * 2 === middleElement) {
21            count++;
22        }
23    }
24  
25    return count;
26}
27

Time and Space Complexity

Time Complexity: O(n), where n is the length of the array nums. The code iterates through the array once using a generator expression within the sum() function. The range is from index 1 to len(nums) - 1, which means we check n - 2 elements. For each element at index i, we perform constant time operations: accessing three array elements (nums[i-1], nums[i], and nums[i+1]), performing arithmetic operations (addition and multiplication), and a comparison. Since all operations inside the loop are O(1) and we iterate through n - 2 elements, the overall time complexity is O(n).

Space Complexity: O(1). The code uses a generator expression rather than creating a list comprehension, which means it doesn't create an intermediate list to store all the boolean values. The sum() function processes the generator expression element by element without storing all results in memory. Only a constant amount of extra space is used for the loop variable i and temporary calculations, regardless of the input size.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Integer Division vs Floating Point Division

A common mistake is directly checking nums[i-1] + nums[i+1] == nums[i] / 2. This can lead to incorrect results when nums[i] is odd, as integer division in some languages truncates the decimal part. For example, with [1, 5, 1], checking 1 + 1 == 5 / 2 would give 2 == 2 in Python 3 (true) but 2 == 2.5 if using float division (false).

Solution: Always multiply both sides by 2 to avoid division: 2 * (nums[i-1] + nums[i+1]) == nums[i]. This ensures we're working with integers throughout and avoids any division-related precision issues.

2. Integer Overflow in Multiplication

When multiplying (nums[i-1] + nums[i+1]) * 2, the result could potentially overflow if the array contains very large integers, especially in languages with fixed integer sizes.

Solution: In languages prone to overflow, consider using appropriate data types (like long in Java/C++) or check for overflow conditions before multiplication. Python handles arbitrary precision integers automatically, so this isn't an issue there.

3. Off-by-One Errors in Loop Boundaries

It's easy to accidentally use range(0, len(nums)) or range(1, len(nums)) instead of range(1, len(nums) - 1), which would cause index out of bounds errors when accessing nums[i+1].

Solution: Always remember that for subarrays of length 3, valid middle positions are from index 1 to len(nums) - 2 inclusive. Double-check loop boundaries and consider adding comments to clarify the reasoning.

4. Misunderstanding the Problem Statement

Some might interpret "half of the middle element" as (a + c) == b * 2 instead of (a + c) == b / 2, reversing the relationship entirely.

Solution: Carefully read the problem statement and test with the provided examples. When [1, 6, 2] satisfies the condition, verify: 1 + 2 = 3 and 6 / 2 = 3, confirming the correct interpretation.

5. Handling Edge Cases with Array Length

Forgetting to handle arrays with length less than 3, which cannot have any valid subarrays of length 3.

Solution: The current implementation naturally handles this since range(1, len(nums) - 1) produces an empty range when len(nums) < 3, resulting in a count of 0. However, adding an explicit check can improve code clarity:

if len(nums) < 3:
    return 0

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan: