What are the usage scenarios of heap and priority queue in Python?

What are the usage scenarios of heap and priority queue in Python?

Heap is a special binary tree structure that is often used to efficiently maintain a dynamic collection. The heapq module in Python provides a heap implementation and can easily perform heap operations.

Priority queue is also a special data structure. Different from ordinary queues, each element of it has a priority associated with it. The highest priority element is taken out first. The heapq module in Python can also implement the priority queue function.

Below we introduce some specific scenarios of using heaps and priority queues, and give relevant code examples.

1. Finding the Top K problem
   It is a common problem to solve the first k largest or smallest elements in a sequence, such as solving the first k largest numbers or the first k smallest numbers. . This problem can be easily solved using a heap of size k or a priority queue.

import heapq

def top_k_smallest(nums, k):
    heap = []
    for num in nums:
        heapq.heappush(heap, num)
        if len(heap) > k:
            heapq.heappop(heap)
    return heap

nums = [5, 3, 8, 2, 7, 1, 9]
k = 3
result = top_k_smallest(nums, k)
print(result)  # 输出 [3, 2, 1]

2. Merge ordered arrays
   It is a common problem to merge multiple ordered numbers to form an ordered array. It can be implemented using a priority queue. Each time, the smallest element is taken from each array and put into the priority queue, and then the elements in the queue are taken out in turn.

import heapq

def merge_sorted_arrays(arrays):
    result = []
    pq = []
    for array in arrays:
        if array:
            heapq.heappush(pq, (array[0], array))
    
    while pq:
        smallest, array = heapq.heappop(pq)
        result.append(smallest)
        if array[1:]:
            heapq.heappush(pq, (array[1], array[1:]))
    
    return result

arrays = [[1, 3, 5], [2, 4, 6], [0, 7, 8]]
result = merge_sorted_arrays(arrays)
print(result)  # 输出 [0, 1, 2, 3, 4, 5, 6, 7, 8]

3. Finding the median
   Finding the median of a sequence is a classic problem. This can be achieved using two heaps, a max heap for the first half of the sequence and a min heap for the second half of the sequence. Keeping the sizes of the two heaps equal or different by one, the median can be taken at the top of the heap.

import heapq

def median(nums):
    min_heap = []
    max_heap = []
    for num in nums:
        if len(max_heap) == 0 or num <= -max_heap[0]:
            heapq.heappush(max_heap, -num)
        else:
            heapq.heappush(min_heap, num)
        
        if len(max_heap) > len(min_heap) + 1:
            heapq.heappush(min_heap, -heapq.heappop(max_heap))
        elif len(min_heap) > len(max_heap):
            heapq.heappush(max_heap, -heapq.heappop(min_heap))
    
    if len(max_heap) > len(min_heap):
        return -max_heap[0]
    elif len(min_heap) > len(max_heap):
        return min_heap[0]
    else:
        return (-max_heap[0] + min_heap[0]) / 2

nums = [4, 2, 5, 7, 1, 8, 3, 6]
result = median(nums)
print(result)  # 输出 4.5

The above are some common usage scenarios and sample codes of heap and priority queue in Python. Heaps and priority queues are some commonly used data structures, and mastering their use is very helpful for solving some complex problems.

The above is the detailed content of What are the usage scenarios of heap and priority queue in Python?. For more information, please follow other related articles on the PHP Chinese website!