Python中的多重集-编程知识-白鹭情

我正在解决 CSES、Traffic Lights上的这个问题：

所以为了有效解决这样的问题，我需要一个类似于串列的 Python 资料结构，但元素的搜索和洗掉需要 O(1) 或更类似于集合的资料结构，但我需要能够插入多个相同的元素并保持顺序。我的问题代码是：

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 =" ")

但是这段代码很慢。在 C 中有一种称为多集的资料结构。但是在 python 中找不到类似的资料结构。任何帮助表示赞赏。

uj5u.com热心网友回复：

您拥有的资料结构lens就像一个多重集，也可以作为Counter. 就时间复杂度而言，算法的瓶颈部分是：

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

这是一个具有线性时间复杂度的操作，因此它使算法成为二次的。

为了改进算法的这一部分，我建议使用堆。有heapq它提供了一个最小堆实作。由于您实际上需要一个最大堆，因此您只需使用负长度即可。

其次，insort还具有线性时间复杂度（尽管使用的时间比max()上面的表达式少）。您可以通过使用自平衡搜索树实作来改进这一点，它没有标准库，但有一些库提供排序串列，例如sortedcontainers.

以下是如何调整代码以实作这两个想法：

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 = " ")