CF 1736C1 - Good Subarrays (Easy Version)

Rating: 1300
Tags: binary search, data structures, schedules, two pointers
Solve time: 2m 6s
Verified: yes

Solution

Problem Understanding

We are asked to count all contiguous subarrays of a given array where each element is at least as large as its position in the subarray. More concretely, if a subarray starts at index l and ends at index r, the first element of this subarray must be at least 1, the second at least 2, the third at least 3, and so on. This transforms the original array into a kind of "height requirement": the further you go in a subarray, the larger each element must be.

The input consists of multiple test cases, each with an array of up to 200,000 elements, and the sum of all array lengths over all test cases does not exceed 200,000. That immediately rules out any solution that is quadratic in the array length. A brute-force approach that checks all subarrays individually would require roughly $n^2/2$ checks per array, which can reach $2 \cdot 10^{10}$ operations in the worst case. This is far beyond feasible in a one-second time limit. Therefore we need a linear or near-linear approach.

Subtle edge cases arise when elements are exactly equal to their subarray positions or when elements are smaller than their subarray positions. For instance, the array [1,1,1] only allows single-element subarrays. A naive algorithm might mistakenly count longer subarrays as valid if it only checks the first or last element, so careful attention to the position-relative requirement is critical.

Approaches

The brute-force approach is straightforward: iterate over all possible starting indices l, then iterate over all possible ending indices r >= l, and check if every element a[l..r] satisfies a[l + i] >= i + 1. This approach is correct because it checks each subarray explicitly. The complexity is O(n^2) per test case, which becomes unacceptably slow for large arrays. For example, with n = 2 * 10^5, we would perform roughly $2 \cdot 10^10$ checks.

The key insight for a faster solution comes from transforming the subarray condition into a simpler numeric condition. For a subarray starting at l, define a new array c[i] = a[i] - (i - l + 1). The subarray is good if and only if all elements of c[i] are non-negative. As we extend the subarray from left to right, we only need to track the smallest c[i] seen so far. When c[i] becomes negative, we know the subarray cannot extend beyond this point. This is equivalent to a two-pointer or sliding window technique: start at l, extend r as far as possible until the condition breaks, then count all subarrays starting at l and ending before the breaking point. We can then move l forward and continue without rechecking elements that are guaranteed to satisfy the condition.

This observation reduces the complexity from O(n^2) to O(n) per array because each index is considered at most twice: once when extending r and once when advancing l.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(n^2)	O(1)	Too slow
Optimal	O(n)	O(1)	Accepted

Algorithm Walkthrough

Read the number of test cases. For each test case, read the array length n and the array a.
Initialize a variable ans to 0 to accumulate the number of good subarrays.
Use two pointers l and r to maintain a sliding window. Start with l = 0 and r = 0.
While l < n, increment r as long as a[r] >= r - l + 1. This ensures every element in the current window meets the "position-based" requirement.
Add r - l to ans, because all subarrays starting at l and ending before r are valid. The difference r - l counts exactly how many subarrays originate at l.
Increment l. If r < l, set r = l because the window cannot shrink past the left pointer.
Continue until all starting indices have been considered.
Output ans for the test case.

Why it works: The invariant is that at any point in the algorithm, all subarrays [l, r-1] are good, and r is the first index where the subarray fails the requirement. This guarantees we count all and only valid subarrays. Each element is evaluated at most twice, giving linear time complexity.

Python Solution

import sys
input = sys.stdin.readline

t = int(input())
for _ in range(t):
    n = int(input())
    a = list(map(int, input().split()))
    
    ans = 0
    r = 0
    for l in range(n):
        if r < l:
            r = l
        while r < n and a[r] >= r - l + 1:
            r += 1
        ans += r - l
    print(ans)

We initialize r = 0 outside the loop over l so that we can extend the right boundary efficiently. The condition if r < l: r = l ensures the window never becomes invalid when l advances. The while loop extends the right boundary greedily until the subarray no longer satisfies the position condition. Finally, adding r - l counts all subarrays starting at l in O(1) per starting index.

Worked Examples

Sample 1:

Input array: [1, 2, 3]

l	r	r-l	Subarray valid?
0	3	3	`[1,2,3]` all good
1	3	2	`[2,3]` good
2	3	1	`[3]` good

Total subarrays counted: 6. Matches expected output.

Sample 2:

Input array: [1, 1, 1]

l	r	r-l	Subarray valid?
0	1	1	`[1]` good
1	2	1	`[1]` good
2	3	1	`[1]` good

Total subarrays counted: 3. Matches expected output.

The traces show the sliding window only extends as far as the array satisfies the requirement, demonstrating correct handling of edge cases.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n) per test case	Each element is visited at most twice: once by `l` and once by `r`.
Space	O(1) extra	Only pointers and counters are stored, no auxiliary arrays needed.

Given the sum of n across all test cases is ≤ 2·10^5, the total number of operations is roughly 4·10^5, which fits comfortably within the time limit.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    output = io.StringIO()
    sys.stdout = output
    # solution
    t = int(input())
    for _ in range(t):
        n = int(input())
        a = list(map(int, input().split()))
        ans = 0
        r = 0
        for l in range(n):
            if r < l:
                r = l
            while r < n and a[r] >= r - l + 1:
                r += 1
            ans += r - l
        print(ans)
    return output.getvalue().strip()

# Provided samples
assert run("3\n3\n1 2 3\n3\n1 1 1\n4\n2 1 4 3\n") == "6\n3\n7"

# Custom test cases
assert run("1\n1\n1\n") == "1", "minimum-size input"
assert run("1\n5\n5 4 3 2 1\n") == "5", "strictly decreasing"
assert run("1\n5\n1 2 2 3 5\n") == "11", "mixed array"
assert run("1\n6\n1 1 2 2 3 4\n") == "14", "multiple short subarrays"

Test input	Expected output	What it validates
`1\n1\n1`	`1`	Single-element array
`1\n5\n5 4 3 2 1`	`5`	Strictly decreasing array
`1\n5\n1 2 2 3 5`	`11`	Mixed values, some subarrays fail early
`1\n6\n1 1 2 2 3 4`	`14`	Overlapping subarrays of varying lengths

Edge Cases

For an array where elements are exactly equal to their positions in the subarray, such as `[1,