Need help with Code Rush X question 2!

Правка en1, от iron_nicko, 2023-01-28 22:03:44

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.

История

 
 
 
 
Правки
 
 
  Rev. Язык Кто Когда Δ Комментарий
en3 Английский iron_nicko 2023-01-28 23:38:58 873 Changed the code
en2 Английский iron_nicko 2023-01-28 22:34:53 50 Spolier-ed the code
en1 Английский iron_nicko 2023-01-28 22:03:44 1681 Initial revision (published)