sweeneyde

Pure-Python non-recursive mergesort-stye algorithm

Dec 19th, 2019
215
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 4.18 KB | None | 0 0
  1. import itertools
  2.  
  3. def multi_merge(*iterables):
  4.     """
  5.    Non-recursive mergesort-style algorithm:
  6.    Maintain a "tree of losers" of comparisons
  7.    as winners advance to the root.
  8.    (c.f. Knuth Volume 3, Chapter 5.4.1. on "Multiway Merging").
  9.  
  10.    Essentially use a heap, but instead of pushing new items to the
  11.    root and sifting up, push them directly to the leaf nodes--items
  12.    only ever move toward the root.
  13.  
  14.    Example for n=5 iterables A--E:
  15.            0
  16.        1       2
  17.      3   4   5   6
  18.     7 8  |   |   |
  19.     | |  C   D   E
  20.     A B
  21.        (shift = -3)
  22.    Always:
  23.        n leaf nodes
  24.            - get their values directly from iterators
  25.        n-1 internal nodes
  26.            - get their values from their children
  27.    """
  28.     n = len(iterables)
  29.     if n == 0:
  30.         return
  31.     if n == 1:
  32.         yield from iterables[0]
  33.         return
  34.  
  35.     tree = [object()] * (n + n - 1)
  36.     sentinel = object()
  37.  
  38.     # For stability, rotate the list so that the first iterator
  39.     # moves from the smallest index to the leftmost index.
  40.     shift = n - (1 << n.bit_length())
  41.     getters =[itertools.chain(iterables[i + shift],
  42.                               itertools.repeat(sentinel)
  43.                               ).__next__
  44.               for i in range(len(iterables))]
  45.  
  46.     # initialize leaves
  47.     for i, getter in enumerate(getters):
  48.         tree[i + n - 1] = getter()
  49.     # initialize internal nodes by drawing from leaves
  50.     for i in reversed(range(n - 1)):
  51.         while i < n - 1:
  52.             left = 2 * i + 1
  53.             right = left + 1
  54.             t_left = tree[left]
  55.             t_right = tree[right]
  56.             if t_right is sentinel:
  57.                 winner = left
  58.             elif t_left is sentinel:
  59.                 winner = right
  60.             elif t_right < t_left:
  61.                 winner = right
  62.             else:
  63.                 winner = left
  64.             tree[i] = tree[winner]
  65.             i = winner
  66.         tree[i] = getters[i - (n - 1)]()
  67.  
  68.     # Main loop:
  69.     #   - yield the overall winner
  70.     #   - replace parents with their better children
  71.     #   - replace leaves using original iterables
  72.     # This is almost the same as repeated heapreplace, except we get the new
  73.     # leaf nodes from the iterator *corresponding to that leaf node* rather
  74.     # rather than from the iterator of the thing we just yielded.
  75.     #     This (1) prevents us having to store many pointers to the iterators
  76.     # (stability comes from the location in the tree) and (2) lifts the
  77.     # requirement that there is exactly one key from each iterator in the
  78.     # heap, making it so that all of the keys in the internal nodes of the
  79.     # heap are actually as similar as possible, reducing unnecessary comparisons.
  80.     while True:
  81.         t0 = tree[0]
  82.         if t0 is sentinel:
  83.             return
  84.         yield t0
  85.         i = 0
  86.         while i < n - 1:
  87.             left = 2 * i + 1
  88.             right = left + 1
  89.             t_left = tree[left]
  90.             t_right = tree[right]
  91.             if t_right is sentinel:
  92.                 winner = left
  93.             elif t_left is sentinel:
  94.                 winner = right
  95.             elif t_right < t_left:
  96.                 winner = right
  97.             else:
  98.                 winner = left
  99.             tree[i] = tree[winner]
  100.             i = winner
  101.         tree[i] = getters[i - (n - 1)]()
  102.  
  103.  
  104. if __name__ == "__main__":
  105.     from timeit import timeit
  106.     import random
  107.     from test import support
  108.  
  109.     py_heapq = support.import_fresh_module('heapq', blocked=['_heapq'])
  110.  
  111.     print("number of iterables, (time to pure-python multi_merge / pure-python heapq merge)")
  112.  
  113.     n = 0
  114.     while True:
  115.         n = max(int(1.2*n), n+1)
  116.         t1 = t2 = 0
  117.         for _ in range(5):
  118.             lists = [
  119.                 sorted(random.random() for _ in
  120.                        range(random.randint(1000, 5000)))
  121.                 for _ in range(n)
  122.             ]
  123.             assert list(multi_merge(*lists)) == list(py_heapq.merge(*lists))
  124.             t1 += timeit(lambda: list(multi_merge(*lists)), number=50)
  125.             t2 += timeit(lambda: list(py_heapq.merge(*lists)), number=50)
  126.         print(n, t1/t2, sep='\t')
Advertisement
Add Comment
Please, Sign In to add comment