3392. Count Subarrays of Length Three With a Condition

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.

Intuition

To solve this problem, the key observation is to recognize each subarray as consisting of three elements: the first, second, and third elements in the subarray. The condition we're checking for is whether twice the sum of the first and third elements equals the second element. This is a direct comparison that can be efficiently checked for each possible subarray of length 3.

We iterate over the array starting from the second element (index 1) to the second last element (since the subarray must consist of three elements). For each element at index i, we consider the subarray formed with elements nums[i-1], nums[i], and nums[i+1]. If the condition (nums[i-1] + nums[i+1]) * 2 == nums[i] is satisfied, we increment our count.

This approach is efficient and works because it leverages the simplicity of the mathematical condition that needs to be checked, allowing us to total up the number of valid subarrays in a single traversal of the array.

Solution Approach

To solve the problem, we employ a single-pass iteration strategy. Here's how the solution is structured:

Iterate Over the Array: We traverse the nums array from the second element to the second-last element. This is done using a for loop with the index i ranging from 1 to len(nums) - 2. This ensures we can check each subarray of length 3.
Check Condition: For each element at index i, evaluate whether the condition (nums[i-1] + nums[i+1]) * 2 == nums[i] holds true. This checks if twice the sum of the first and third elements of the subarray is equal to the second element.
Count Valid Subarrays: Use the expression sum((nums[i-1] + nums[i+1]) * 2 == nums[i] for i in range(1, len(nums) - 1)) to compute the total number of valid subarrays. This sum comprehension will iterate over the valid indices, evaluate the condition, and sum the boolean results where True is equivalent to 1.
Return Result: The computed sum is the final count of the subarrays that meet the specified condition, and it is returned as the function's output.

The solution effectively uses a linear scan with constant space, making it optimal for the problem size constraints typically encountered in such problems.

Ready to land your dream job?

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

Start Evaluator

Example Walkthrough

Let's take a small example to illustrate the solution approach:

Consider the array nums = [4, 8, 6, 4, 10]. We want to find the number of subarrays of length 3 such that the sum of the first and third numbers equals exactly half of the second number.

Iterate Over the Array:

We start iterating from the second element and go up to the second last element. In this array, we will consider the indices 1, 2, and 3 for forming subarrays.
Subarray Checking:
- For i = 1, the subarray is [4, 8, 6].
  - Check if (4 + 6) * 2 == 8. Here it calculates to 20 == 8, which is false.
- For i = 2, the subarray is [8, 6, 4].
  - Check if (8 + 4) * 2 == 6. Here it calculates to 24 == 6, which is false.
- For i = 3, the subarray is [6, 4, 10].
  - Check if (6 + 10) * 2 == 4. Here it calculates to 32 == 4, which is false.
Count Valid Subarrays:

None of the subarrays satisfy the condition, so the count of valid subarrays is 0.
Return Result:

The function returns 0 since no valid subarrays were found.

This example demonstrates how we smoothly iterate through potential subarrays and apply the condition to check for validity using the given formula.

Solution Implementation

1from typing import List
2
3class Solution:
4    def countSubarrays(self, nums: List[int]) -> int:
5        # Initialize a sum to count the number of subarrays meeting the condition
6        return sum(
7            # Check if the current element (nums[i]) is equal to double of the average of its neighbors (nums[i-1] and nums[i+1])
8            (nums[i - 1] + nums[i + 1]) * 2 == nums[i]
9            for i in range(1, len(nums) - 1)
10        )
11
12# The above list comprehension counts the number of subarrays where the middle element is twice the average of its neighbors.
13

1class Solution {
2    public int countSubarrays(int[] nums) {
3        int count = 0; // Initialize the count of subarrays
4      
5        // Iterate through the array from the second to the second last element
6        for (int i = 1; i + 1 < nums.length; ++i) {
7            // Check if the middle element is the average of its adjacent elements
8            if ((nums[i - 1] + nums[i + 1]) * 2 == nums[i]) {
9                ++count; // Increment count if the condition is true
10            }
11        }
12      
13        return count; // Return the total count
14    }
15}
16

1#include <vector>
2
3class Solution {
4public:
5    // Function to count the number of subarrays that satisfy a certain condition
6    int countSubarrays(std::vector<int>& nums) {
7        int ans = 0; // Initialize the answer to 0
8        // Loop through the array starting from the second element and ending at the second last
9        for (int i = 1; i + 1 < nums.size(); ++i) {
10            // Check if the middle element is equal to the average of its neighbors
11            if ((nums[i - 1] + nums[i + 1]) * 2 == nums[i]) {
12                ++ans; // If the condition is satisfied, increment the answer
13            }
14        }
15        return ans; // Return the total count
16    }
17};
18

1/**
2 * Counts the number of special subarrays in the given array.
3 * A subarray is considered special if it follows the condition:
4 * (nums[i - 1] + nums[i + 1]) * 2 === nums[i]
5 *
6 * @param nums - The input array of numbers.
7 * @returns The count of special subarrays.
8 */
9function countSubarrays(nums: number[]): number {
10    let count: number = 0; // Initialize the count of special subarrays
11  
12    // Iterate through the array from the second element to the second to last element
13    for (let index = 1; index + 1 < nums.length; ++index) {
14        // Check if the current element satisfies the special condition
15        if ((nums[index - 1] + nums[index + 1]) * 2 === nums[index]) {
16            ++count; // Increment the count if the condition is met
17        }
18    }
19  
20    return count; // Return the total count
21}
22

Time and Space Complexity

The time complexity of the code is O(n), where n is the length of the array nums, because the code iterates over each element exactly once (with the exception of the first and last elements). The sum function processes each boolean expression, making it proportional to the size of the list.

The space complexity is O(1) since the solution only uses a fixed amount of additional space for the computation, regardless of the input size. The space used depends only on the variables needed for the sum and loop iteration.

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

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

Question 1 out of 10

Which of the following array represent a max heap?

20 12 16 1 2 3 4

20 30 10 5 6 7 18

30 20 11 2 3 4 12

50 30 40 10 20 30 40

I don’t know