-
Notifications
You must be signed in to change notification settings - Fork 2
/
imerge.py
43 lines (34 loc) · 1.15 KB
/
imerge.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
import heapq
def imerge(*iterables):
'''Merge multiple sorted inputs into a single sorted output.
Equivalent to: sorted(itertools.chain(*iterables))
>>> list(imerge([1,3,5,7], [0,2,4,8], [5,10,15,20], [], [25]))
[0, 1, 2, 3, 4, 5, 5, 7, 8, 10, 15, 20, 25]
'''
heappop, siftup, _StopIteration = heapq.heappop, heapq._siftup, StopIteration
h = []
h_append = h.append
for it in map(iter, iterables):
try:
next = it.next
h_append([next(), next])
except _StopIteration:
pass
heapq.heapify(h)
last = None
while 1:
try:
while 1:
v, next = s = h[0] # raises IndexError when h is empty
if v[0] == last[0]:
last
yield v
s[0] = next() # raises StopIteration when exhausted
siftup(h, 0) # restore heap condition
except _StopIteration:
heappop(h) # remove empty iterator
except IndexError:
return
if __name__ == '__main__':
import doctest
doctest.testmod()