大きい二つのリストをマージする

たまに、二つのソート済みリストを大きい一つのソート済みリストにしたいときがある。

具体的にはこんな状況:

>>> a = [1, 3, 6, 8, 12]
>>> b = [2, 4, 5, 9, 10]
>>> sorted(a + b)
[1, 2, 3, 4, 5, 6, 8, 9, 10, 12]

2つのリストaとbが比較的小さいリストならいいけど、
例えばイテレータ・ジェネレータオブジェクトだったり巨大なリストだったりするときに
役に立つのがheapq。

>>> import heap

各々のリストはソートしてある必要があるけど、
2つのリストをマージしてくれる。

>>> import itertools
>>> a = itertools.takewhile(lambda x: x < 10 ** 7, (5 * x for x in itertools.count(1)))  # range(5, 10 ** 7, 5) と一緒
>>> b = itertools.takewhile(lambda x: x < 10 ** 7, (7 * x for x in itertools.count(1)))  # range(7, 10 ** 7, 7) と一緒
>>> list(itertools.islice(heapq.merge(a, b), 100))
[5, 7, 10, 14, 15, 20, 21, 25, 28, 30, 35, 35, 40, 42, 45, 49, 50, 55, 56, 60, 63, 65, 70, 70, 75, 77, 80, 84, 85, 90, 91, 95, 98, 100, 105, 105, 110, 112, 115, 119, 120, 125, 126, 130, 133, 135, 140, 140, 145, 147, 150, 154, 155, 160, 161, 165, 168, 170, 175, 175, 180, 182, 185, 189, 190, 195, 196, 200, 203, 205, 210, 210, 215, 217, 220, 224, 225, 230, 231, 235, 238, 240, 245, 245, 250, 252, 255, 259, 260, 265, 266, 270, 273, 275, 280, 280, 285, 287, 290, 294]

マージソートとかに大変役に立つ。あとは無駄にリストを作りたくない場合とか。

ついでに、groupbyを使えばuniqも同時にかけられる。

>>> list(itertools.islice((x for x, _ in itertools.groupby(heapq.merge(a, b))), 100))
[5, 7, 10, 14, 15, 20, 21, 25, 28, 30, 35, 40, 42, 45, 49, 50, 55, 56, 60, 63, 65, 70, 75, 77, 80, 84, 85, 90, 91, 95, 98, 100, 105, 110, 112, 115, 119, 120, 125, 126, 130, 133, 135, 140, 145, 147, 150, 154, 155, 160, 161, 165, 168, 170, 175, 180, 182, 185, 189, 190, 195, 196, 200, 203, 205, 210, 215, 217, 220, 224, 225, 230, 231, 235, 238, 240, 245, 250, 252, 255, 259, 260, 265, 266, 270, 273, 275, 280, 285, 287, 290, 294, 295, 300, 301, 305, 308, 310, 315, 320]

2つあった35に注目。

ということで、itertoolsも便利だなと。