Merge sorted iterables

imerge_sorted.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright © 2015-present Antoine Catton
# This work is free. You can redistribute it and/or modify it under the
# terms of the Do What The Fuck You Want To Public License, Version 2,
# as published by Sam Hocevar. See http://www.wtfpl.net/ for more details

def imerge_sorted(iterables, key=lambda x: x, reversed=False):
    """
    https://gist.github.com/acatton/b362585d501209d293d5

    Merge sorted iterables::

        >>> list(imerge_sorted([
        ...     ['A', 'E', 'F'],
        ...     ['B', 'C', 'H'],
        ...     ['D', 'G', 'I'],
        ... ]))
        ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I']

        >>> list(imerge_sorted([
        ...     ['Z', 'Y', 'U'],
        ...     ['X', 'V', 'T'],
        ...     ['W', 'S', 'R'],
        ... ], reversed=True))
        ['Z', 'Y', 'X', 'W', 'V', 'U', 'T', 'S', 'R']

     * If each iterable hasn't been sorted, the result is unpredictable.
     * If two elements are equal, the result could be unstable.
     * The result is an iterator.
     * This is like heapq.merge except that you can provide a key function and
       ask for reversed sorting.

    For merging n iterables each of them of size m, this function would take:
     * A worst case computational complexity of O(n * m)
     * A worst case space complexity of O(n)

    On the other hand, heapq.merge has a worst space complexity of O(n * m)...
    """

    def select_next_element(items, key, reversed):
        """
        Return the next element to take from the items dict.
        """
        selection_func = max if reversed else min
        key_func = lambda x: key(x[1])
        return selection_func(items.items(), key=key_func)

    row = tuple(iter(i) for i in iterables)
    items = {n: next(i) for n, i in enumerate(row)}

    while len(items) > 0:
        n, elem = select_next_elem(items, key, reversed)
        yield elem

        try:
            items[n] = next(row[n])
        except StopIteration:
            items.pop(n)