76 lines
1.8 KiB
Python
76 lines
1.8 KiB
Python
import random
|
|
from k_way_merge import merge_k_sorted, merge_k_sorted_list
|
|
|
|
|
|
def test_empty():
|
|
assert merge_k_sorted_list([]) == []
|
|
assert merge_k_sorted_list([[], [], []]) == []
|
|
|
|
|
|
def test_single():
|
|
assert merge_k_sorted_list([[1, 2, 3]]) == [1, 2, 3]
|
|
|
|
|
|
def test_basic():
|
|
a = [1, 4, 7]
|
|
b = [2, 5, 8]
|
|
c = [3, 6, 9]
|
|
assert merge_k_sorted_list([a, b, c]) == list(range(1, 10))
|
|
|
|
|
|
def test_duplicates_stability():
|
|
# equal values across streams: lower index first
|
|
a = [5, 5]
|
|
b = [5, 5]
|
|
out = merge_k_sorted_list([a, b])
|
|
assert out == [5, 5, 5, 5]
|
|
|
|
|
|
def test_empty_streams_mixed():
|
|
assert merge_k_sorted_list([[], [1], [], [2, 3]]) == [1, 2, 3]
|
|
|
|
|
|
def test_iterables_not_lists():
|
|
def gen(n):
|
|
for i in range(n):
|
|
yield i * 2
|
|
assert merge_k_sorted_list([gen(3), [1, 3, 5]]) == [0, 1, 2, 3, 4, 5]
|
|
|
|
|
|
def test_large_random_vs_sorted():
|
|
rng = random.Random(42)
|
|
streams = []
|
|
for _ in range(7):
|
|
s = sorted(rng.randint(0, 1000) for _ in range(200))
|
|
streams.append(s)
|
|
out = merge_k_sorted_list(streams)
|
|
expected = sorted(x for s in streams for x in s)
|
|
assert out == expected
|
|
|
|
|
|
def test_total_count():
|
|
streams = [[1, 2], [3], [], [4, 5, 6]]
|
|
assert len(merge_k_sorted_list(streams)) == 6
|
|
|
|
|
|
def test_single_elements():
|
|
assert merge_k_sorted_list([[9], [1], [5]]) == [1, 5, 9]
|
|
|
|
|
|
def test_descending_not_assumed():
|
|
# inputs assumed sorted ascending; verify we don't sort inputs
|
|
assert merge_k_sorted_list([[1, 2], [3, 4]]) == [1, 2, 3, 4]
|
|
|
|
|
|
def test_iterator_lazy():
|
|
it = merge_k_sorted([[1, 2], [3]])
|
|
assert next(it) == 1
|
|
assert next(it) == 2
|
|
assert next(it) == 3
|
|
try:
|
|
next(it)
|
|
except StopIteration:
|
|
pass
|
|
else:
|
|
raise AssertionError("expected StopIteration")
|