Problem Informations

• Problem Index: 3392
• Problem Link: Count Subarrays of Length Three With a Condition
• Topics: Array

Problem Description

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 \leq \text{nums.length} \leq 100$
• $-100 \leq \text{nums}[i] \leq 100$

Intuition

The problem requires counting subarrays of length 3 where the sum of the first and third element equals exactly half of the middle element. Recognizing that a brute force solution is feasible due to the constraints, the primary intuition is to iterate over possible subarrays, checking the specified condition. The problem imposes that the subarray has exactly three elements, simplifying the approach.

Approach

We start by initializing a counter variable to zero, which will hold the number of valid subarrays. We then iterate through the array with an index i such that the subarray from index i to i+2 is considered. The looping stops at length - 2 to prevent accessing indices beyond the array’s boundary.

For each subarray nums[i], nums[i+1], nums[i+2], we check the condition: $$2 \times (nums[i] + nums[i+2]) = nums[i+1]$$. This condition ensures that the sum of the first and third elements equals half of the middle element. If the condition is satisfied, the counter is incremented. After iterating through the array, the counter, which now holds the number of valid subarrays, is returned.

This approach efficiently handles the constraints given, iterating through the list in linear time, and checking conditions on each triplet of numbers.

Code

C++

class Solution {
public:
    int countSubarrays(vector<int>& nums) {
        int count = 0;
        int length = nums.size();

        // Iterate over the array, considering each set of 3 consecutive elements
        for (int i = 0; i < length - 2; i++) {
            // Check if the sum of the first and third numbers is equal to half of the second number
            if (2 * (nums[i] + nums[i + 2]) == nums[i + 1]) {
                count++; // Increment the count if the condition is satisfied
            }
        }
        return count; // Return the total count of valid subarrays
    }
};

Python

class Solution:
    def countSubarrays(self, nums):
        count = 0
        length = len(nums)

        # Iterate over the array, considering each set of 3 consecutive elements
        for i in range(length - 2):
            # Check if the sum of the first and third numbers is equal to half of the second number
            if 2 * (nums[i] + nums[i + 2]) == nums[i + 1]:
                count += 1  # Increment the count if the condition is satisfied
        
        return count  # Return the total count of valid subarrays

Complexity

• Time complexity: The time complexity of the given algorithm is $O(n)$, where $n$ is the length of the input array nums. This is because the algorithm iterates through the array a single time, checking each possible subarray of length 3 exactly once.

• Space complexity: The space complexity of the algorithm is $O(1)$. The algorithm uses a constant amount of additional space, regardless of the input size, for variables like count and length.