Multisets in Python

Question

I am working on this problem on CSES, Traffic Lights:

There is a street of length 𝑥 whose positions are numbered 0,1,…,𝑥. Initially there are no traffic lights, but 𝑛 sets of traffic lights are added to the street one after another.

Your task is to calculate the length of the longest passage without traffic lights after each addition.

Input

The first input line contains two integers 𝑥 and 𝑛: the length of the street and the number of sets of traffic lights.

Then, the next line contains n integers 𝑝₁,𝑝₂,…,𝑝_𝑛: the position of each set of traffic lights. Each position is distinct.

Output

Print the length of the longest passage without traffic lights after each addition.

Constraints

• 1 ≤ 𝑥 ≤ 10⁹
• 1 ≤ 𝑛 ≤ 2⋅10⁵
• 0 < 𝑝_𝑖 < 𝑥

Example

Input:

8 3
3 6 2

Output:

5 3 3

So to effectively solve a problem like this, I need a data structure in Python similar to a list but the search and deletion of elements need to be O(1) or more like a data structure similar to sets but I need to be able to insert multiple same elements and also preserve order. My code for the problem is:

from collections import defaultdict
from bisect import bisect_right , insort
x , n = list(map(int  , input().split()))
arr = list(map(int , input().split()))
lens = defaultdict(int)
lens[x] = 1
lights = [0,x]
for ele in arr:
    idx = bisect_right(lights , ele)
    to_be_removed = lights[idx] - lights[idx-1]
    lens[to_be_removed] -= 1
    lens[lights[idx]-ele] += 1
    lens[ele-lights[idx-1]] += 1
    insort(lights , ele)
    print(max([x for x in lens.keys() if lens[x]])  , end =" ") 

However this code is slow. There is a data structure called multi-sets in c++. However couldn't find a similar data structure in python. Any help appreciated.

Answer 1

The data structure you have for lens is like a multiset, also available as Counter. The part of your algorithm that is the bottle neck in terms of time complexity, is this:

max([x for x in lens.keys() if lens[x]]) 

This is an operation with linear time complexity, and so it makes the algorithm quadratic.

To improve on that part of the algorithm, I'd suggest using a heap. There is heapq which provides a min heap implementation. As you actually need a max heap, you'd just feed it with negative lengths.

Secondly, insort also has a linear time complexity (although using less time than the max() expression above). You can improve on this by using a self-balancing search tree implementation, for which there is no standard library, but there are libraries that provide sorted lists such as sortedcontainers.

Here is how you can adjust your code to implement those two ideas:

from collections import defaultdict
from heapq import heappush, heappop
from sortedcontainers import SortedList

x , n = list(map(int  , input().split()))
arr = list(map(int , input().split()))

lens = defaultdict(int)
lens[x] = 1
lights = SortedList([0, x])  # For faster insertion
heap = [-x]  # Put total width also in a heap
for ele in arr:
    idx = lights.bisect_right(ele)
    to_be_removed = lights[idx] - lights[idx-1]
    lens[to_be_removed] -= 1

    # Add widths to the heap when they are the only occurrences
    right = lights[idx]-ele
    if lens[right] == 0:
        heappush(heap, -right)
    lens[right] += 1

    left = ele-lights[idx-1]
    if lens[left] == 0:
        heappush(heap, -left)
    lens[left] += 1

    # Remove the largest width as long as it no longer represents a segment
    while lens[-heap[0]] == 0:
        heappop(heap)
    
    # The add method is O(logn)
    lights.add(ele)
    # Just output the largest width in the heap
    print(-heap[0], end = " ")