USACO Bronze 2023 Open - FEB

Debugging Help
#1

Checklist

Make sure to read all of How to ask for help on a problem before asking for help.

Problem Info

Name: FEB
Link: Solution - FEB (USACO Bronze 2023 Open)

Question

I spent a long time working out all the cases on paper and came to a solution that is correct but too slow. I think the time complexity is O(N) so I assumed that it would work for N <= 2 * 10 ** 5, but it clearly doesn’t as it fails all the 2*10**5 test cases (test data). I would mainly like to know why the code is too slow and how I could possibly still use the same logic and solve the problem fully.

What I’ve Tried

I wrote a brute force script checking all 2**n possible combinations that also passes the first 8 test cases, and I did a bit of stress testing with small n (<= 10). I am pretty confident that my solution will not get wrong answer on any cases.

My Work

n = int(input())
s = input()
ans = [0]
# ans array is modified to account for the new possibilities that each segment brings. 
# Ex: [0,1] + [1] = [1,2] ||| [0,1] + [0,2] = [0,1,2,3]

#now get all the segments --> F...FE, EF...F, EF...FE, EF...FB etc.
segments = []
currsegment = []
for i in range(n):
    if i > 0 and s[i] == s[i-1] and s[i] != "F":
        ans[0] += 1
    if not currsegment or s[i] == "F":
        currsegment.append(s[i])
        continue
    if s[i] != "F":
        if currsegment[-1] == "F":
            currsegment.append(s[i])
            segments.append(currsegment)
        currsegment = [s[i]]            
segments.append(currsegment)

#find the case of each segment and update answer based on the pattern
for segment in segments:
    length = len(segment)
    if segment[0] == "F" or segment[-1] == "F":
        # F...FE, EF...F ----> length-1,length-2,length-3 etc.
        excitement = [x for x in range(length)]
    elif segment[0] == segment[-1]:
        # EF...FE, EF...FE ----> length-1,length-3,length-5 etc.
        excitement = [x for x in range(length-1,-1,-2)][::-1]
    else:
        # EF...FB, BF...FE ----> length-2,length-4,length-6 etc.
        excitement = [x for x in range(length-2,-1,-2)][::-1]
    new_ans = set()
    #this double loop to update might be the reason it gets TLE
    for elem in ans:
        for addition in excitement:
            new_ans.add(elem + addition)
    ans = list(new_ans)

print(len(ans))
for val in ans:
    print(val)
#2

Solved