Search code examples
pythonheappriority-queue

Python heapq: Split and merge into a ordered heapq


I am wanting to split two heapqs (used as a priority queues), and then add them together and have the resulting heapq ordered in relation to both of the previous heapqs.

Is this possible in python?

My current code:

population = []
for i in range(0, 6):
    heappush(population, i)
new_population = []
for i in range(4, 9):
    heappush(new_population, i)

split_index = len(population) // 2
temp_population = population[:split_index]
population = new_population[:split_index] + temp_population
print(population)
print(heappop(population))

Output:

[4, 5, 6, 0, 1, 2]
4

Wanted output:

[0, 1, 2, 4, 5, 6]
0

Solution

  • Use nlargest instead of slicing, then reheapify the combined lists.

    from heapq import nlargest, heapify
    n = len(population) // 2
    population = heapify(nlargest(population, n) +
                         nlargest(new_population, n))
    print(heappop(population))
    

    You may want to benchmark, though, if sorting the two original lists, then merging the results, is faster. Python's sort routine is fast for nearly sorted lists, and this may impose a lot less overhead than the heapq functions. The last heapify step may not be necessary if you don't actually need a priority queue (since you are sorting them anyway).

    from itertools import islice
    from heapq import merge, heapify
    n = len(population)  # == len(new_population), presumably
    population = heapify(islice(merge(sorted(population), sorted(new_population)), n))