diff --git a/Lib/heapq.py b/Lib/heapq.py index fabefd87f8bf8c..78dd527c979e3e 100644 --- a/Lib/heapq.py +++ b/Lib/heapq.py @@ -329,6 +329,111 @@ def merge(*iterables, key=None, reverse=False): ''' + if len(iterables) == 2: + # special case for performance + + a_iter = iter(iterables[0]) + b_iter = iter(iterables[1]) + next_a = a_iter.__next__ + next_b = b_iter.__next__ + + try: + a = next_a() + except StopIteration: + yield from b_iter + return + try: + b = next_b() + except StopIteration: + yield a + yield from a_iter + return + + if key is None: + if reverse: + # no key, reverse + while True: + if a < b: + yield b + try: + b = next_b() + except StopIteration: + yield a + yield from a_iter + return + else: + yield a + try: + a = next_a() + except StopIteration: + yield b + yield from b_iter + return + else: + # no key, forward + while True: + if b < a: + yield b + try: + b = next_b() + except StopIteration: + yield a + yield from a_iter + return + else: + yield a + try: + a = next_a() + except StopIteration: + yield b + yield from b_iter + return + else: + ka = key(a) + kb = key(b) + if reverse: + # using a key, reverse + while True: + if ka < kb: + yield b + try: + b = next_b() + except StopIteration: + yield a + yield from a_iter + return + kb = key(b) + else: + yield a + try: + a = next_a() + except StopIteration: + yield b + yield from b_iter + return + ka = key(a) + else: + # using a key, forward + while True: + if kb < ka: + yield b + try: + b = next_b() + except StopIteration: + yield a + yield from a_iter + return + kb = key(b) + else: + yield a + try: + a = next_a() + except StopIteration: + yield b + yield from b_iter + return + ka = key(a) + h = [] h_append = h.append diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index 861ba7540df2bc..bf786e37fd0b29 100644 --- a/Lib/test/test_heapq.py +++ b/Lib/test/test_heapq.py @@ -176,22 +176,23 @@ def test_heapsort(self): self.assertEqual(heap_sorted, sorted(data)) def test_merge(self): - inputs = [] - for i in range(random.randrange(25)): - row = [] - for j in range(random.randrange(100)): - tup = random.choice('ABC'), random.randrange(-500, 500) - row.append(tup) - inputs.append(row) - - for key in [None, itemgetter(0), itemgetter(1), itemgetter(1, 0)]: - for reverse in [False, True]: - seqs = [] - for seq in inputs: - seqs.append(sorted(seq, key=key, reverse=reverse)) - self.assertEqual(sorted(chain(*inputs), key=key, reverse=reverse), - list(self.module.merge(*seqs, key=key, reverse=reverse))) - self.assertEqual(list(self.module.merge()), []) + for n in range(26): + inputs = [] + for i in range(n): + row = [] + for j in range(random.randrange(100)): + tup = random.choice('ABC'), random.randrange(-500, 500) + row.append(tup) + inputs.append(row) + + for key in [None, itemgetter(0), itemgetter(1), itemgetter(1, 0)]: + for reverse in [False, True]: + seqs = [] + for seq in inputs: + seqs.append(sorted(seq, key=key, reverse=reverse)) + self.assertEqual(sorted(chain(*inputs), key=key, reverse=reverse), + list(self.module.merge(*seqs, key=key, reverse=reverse))) + self.assertEqual(list(self.module.merge()), []) def test_empty_merges(self): # Merging two empty lists (with or without a key) should produce diff --git a/Misc/NEWS.d/next/Library/2019-12-01-04-01-32.bpo-38938.Mj9nKR.rst b/Misc/NEWS.d/next/Library/2019-12-01-04-01-32.bpo-38938.Mj9nKR.rst new file mode 100644 index 00000000000000..88fbeec7e16759 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2019-12-01-04-01-32.bpo-38938.Mj9nKR.rst @@ -0,0 +1 @@ +For performance, added a special case to heapq.merge for two iterables. \ No newline at end of file