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:
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.