I tried solving Traffic lights problem in the CSES problem set, my approach seems to work for small inputs but I'm facing TLE for large inputs

user_9844638 0
Viewed 856

This is the question. I'm maintaining two lists, one to store the current possible lengths (list 1 i.e., "len"), another (list 2 i.e., "street") to store the street locations with lights in sorted order (along with starting and endpoints).

My approach is to find the lower bound (just less than the current input) and upper bound for input and add the new lengths of passage without any light and delete the old passage distance between lower and upper bound from the list 1, then insert the input at the appropriate position in the list 2. I'm then printing the max of list 1 (lengths).

Code:

import bisect

x, n = [int(x) for x in input().split()] 

lights = list(map(int, input().split()))

# Lists
len, street = [x], [0, x]

i = 0 
while(i < n):
    r = bisect.bisect_right(street, lights[i]) #upper bound
    l = r-1 #lower bound

    len.remove(street[r] - street[l])

    len.append(lights[i] - street[l])
    len.append(street[r] - lights[i])

    street.insert(r, lights[i])

    print(max(len), end = " ")
    i+= 1

I'm getting TLE for large inputs. I feel the "street.insert(r, lights[i])" part of the code is the reason for the error, but, my approach is based on having a sorted list (so that bisect function works to find the lb, ub), so how can I optimise it? Could anyone help me to solve the error?

python-3.x search sorting algorithm
2 Answers

Nice approach. I think your idea is correct. You should know though that removing from a list is an O(N) operation.

len.remove(street[r] - street[l])

if you can swap your list for a dictionary, deleting a key will be O(1) and your solution should not time out on larger input sets.

user_82900600

insert() and remove() on lists are O(n) operations, you can instead replace len in your code with a dictionary. Here is my code with dict; but it still gave a TLE.

import bisect 

I = lambda : map(int, input().split())

n, x = I()

bes = {} # similar to len in your code, but as dict for O(1) update

arr = [0, n] # similar to street
bes[0] = n

for p in I():
    pos = bisect.bisect(arr, p)
    bes[arr[pos - 1]] = p - arr[pos - 1]
    bes[p] = arr[pos] - p
    arr = arr[:pos] + [p] + arr[pos:]
    print(max(bes.values()))

Note that max() operation has an O(n) complexity too. We can avoid max() in every iteration by storing the argmax() and calling max() only when there is an update on argmax(). But it seems that even a bisect.bisect() operation is problematic. I submitted a simpler code with only bisect() and it gave a TLE. I used PyPy3 for compilation.

...
for p in I():
    pos = bisect.bisect(arr, p)
    arr = arr[:pos] + [p] + arr[pos:]

print(1)

There are some other problems in CSES that give TLE for Python (it's terribly slow!) despite using the right algorithm. CSES Hacking section has some great Pythonic solutions for these problems, but you'll have to solve the question first. I typically try out the same algorithm in C++ and then check the Python solutions in Hacking section. Here's one such solution (via zdu863) that completely avoids bisect

I = lambda : map(int, input().split())

n, x = I()

P = list(I())
arr = [0] + sorted(P) + [n, 0] # [0] added for empty slot

l = len(arr) - 1
adj = {arr[i]: [arr[i - 1], arr[i + 1]] for i in range(l)}

max_val = max(arr[i] - arr[i-1] for i in range(1, l))
bes = [max_val]
# note that the current state has all lights kept
# lets remove one by one
for p in P[::-1]:
    left, right = adj.pop(p)
    max_val = max(right - left, max_val) # if light at p was not there
    adj[left][1] = right
    adj[right][0] = left
    bes.append(max_val)

# print(adj)
print(*bes[::-1][1:]) # reverse and print
user_109856500
Related