Need help with Code Rush X question 2!

#	User	Rating
1	tourist	3985
2	jiangly	3814
3	jqdai0815	3682
4	Benq	3529
5	orzdevinwang	3526
6	ksun48	3517
7	Radewoosh	3410
8	hos.lyric	3399
9	ecnerwala	3392
9	Um_nik	3392

#	User	Contrib.
1	cry	169
2	maomao90	162
2	Um_nik	162
4	atcoder_official	161
5	djm03178	158
6	-is-this-fft-	157
7	adamant	155
8	awoo	154
8	Dominater069	154
10	luogu_official	151

Hey, I know the contest was postponed, but I was trying to solve this question:

Question:

Consider the integer sequence A[], where the elements are first N natural numbers in order.

You are now given two integers, L and S. Determine whether there exists a subarray with length L and sum S after removing at most one element from A.

A subarray of an array is a non-empty sequence obtained by removing zero or more elements from the front of the array, and zero or more elements from the back of the array.
1 <= N <= 10^9
1 <= L <= N - 1

I came up with something like this:

My Code


def solve(n, l, s):
    lo = 1
    hi = n

    def helper(i):
        one = i * (i+1) >> 1
        two = 0
        if i - l > 0:
            two = (i-l)*(i-l+1) >> 1
        return one - two
    for _ in range(100):
        mid = ((hi - lo) >> 1) + lo
        val = helper(mid)
        if abs(s - val) < 2:
            if lo == 1:
                if s - val >= 0:
                    return "YES"
            elif mid == n:
                if s - val <= 0:
                    return "YES"
            else:
                return "YES"
        if s > val:
            lo = mid + 1
        elif s < val:
            hi = mid - 1
    return "NO"


for _ in range(int(input())):
    n, l, s = map(int, input().split())

    print(solve(n, l, s))

Test Cases with the answers:

3
5 3 11 # YES
5 3 5 # NO
5 3 6 # YES

My code seems to pass the initial test cases, but I'd like to know if this will work for all cases. Please let me know if this is correct.

Updated Code

# Author : iron_nicko
from collections import Counter, defaultdict, deque
from bisect import bisect_left, bisect_right
from functools import cache
from itertools import accumulate
from heapq import heapify, heappush, heappop


def range_sum(l, r):
    l -= 1
    x = l * (l + 1) >> 1
    y = r * (r + 1) >> 1
    return y - x

def check(N, L, S, deletion=0):
    L += deletion
    l, r = 1,  N - L + 1
    while l <= r:
        mid = (l + r) >> 1
        subarray_sum = range_sum(mid, mid + L - 1)
        diff = subarray_sum - S
        if not deletion:
            if subarray_sum == S:
                return True
            elif subarray_sum < S:
                l = mid + 1
            else:
                r = mid - 1
        else:
            if diff < mid:
                l = mid + 1
            elif mid <= diff <= mid + L - 1:
                return True
            else:
                r = mid - 1
    return False


def solve():
    N, L, S = map(int, input().split())
    print("NYOE S"[check(N, L, S, 1) or check(N, L, S) :: 2].rstrip())

for _ in range(int(input())):
    solve()

iron_nicko's blog