Need help with Code Rush X question 2!

№	Пользователь	Рейтинг
1	tourist	3845
2	jiangly	3707
3	Benq	3630
4	orzdevinwang	3573
5	Geothermal	3569
5	cnnfls_csy	3569
7	jqdai0815	3532
8	ecnerwala	3501
9	gyh20	3447
10	Rebelz	3409

№	Пользователь	Вклад
1	maomao90	171
2	awoo	163
2	adamant	163
4	maroonrk	152
5	nor	151
5	-is-this-fft-	151
7	TheScrasse	147
7	atcoder_official	147
9	Petr	145
10	pajenegod	144

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

def solve(n, l, s):
    lo = l
    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):
        if hi < lo: break
        mid = ((hi - lo) >> 1) + lo
        val = helper(mid)
        if abs(s - val) < 2:
            if lo == l:
                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))

Блог пользователя iron_nicko