Selection rank algorithm demo

selrank.py

"""Selecton rank algorithm
"""

import random

def selrank(array, start, stop, k):
    """return the k-th largest element (0-th = max) in sliced list
    `array[start:stop]` using the selection rank algorithm

    This move the larger elements of array to lower indices such that at the end
    array[k] is the result to return
    """
    print("array(%s) start(%s) stop(%s) k(%s)" % (array, start, stop, k))
    assert stop > start, "Range error in start:stop"
    assert k >= 0, "k must not be negative"
    assert k <= (stop-start-1), "k too small for range start:stop"
    assert start >= 0, "start must be non-negative"
    assert len(array) >= stop, "stop too small for input array"
    # corner case
    if stop-start-1 == 0:
        return array[start]
    # pick pivot randomly
    pivot_idx = random.randint(start, stop-1)
    pivot_val = array[pivot_idx]
    print("pivot_idx(%s) pivot_val(%s)" % (pivot_idx, pivot_val))
    cursor = start  # pos to place next item that is greater than or equal to pivot_val
                    # invariant: always assert all(x >= pivot_val for x in array[start:cursor])
    for i in range(cursor, stop):
        if array[i] >= pivot_val:
            array[i], array[cursor] = array[cursor], array[i]
            cursor += 1
    print("final cursor(%s)" % cursor)
    # check and return: there is a chance that we picked the min as pivot and
    # cursor == stop at this point, in which the current recursion does nothing
    # to narrow down the search scope. If we change the criteria to
    # array[i] > pivot_val above, we may also picked the max as pivot and
    # similarly does nothing. Therefore picking a random index is important to
    # avoid infinite loop.
    if cursor - start == k + 1:
        # len(array[start:cursor]) == k+1, pivot_val is the one
        return pivot_val
    elif cursor - start > k + 1:
        # k-th larget is somewhere at lower indices than cursor
        return selrank(array, start, cursor, k)
    else:
        # k-th larget is somewhere at higher or equal indices than cursor
        return selrank(array, cursor, stop, k-(cursor-start))

if __name__ == "__main__":
    LIST = [3, 14, 1, 59, 26, 53, 58, 97, 93, 23, 84, 62, 64, 33]
    print("List to begin: %s" % LIST)
    v = selrank(LIST, 0, len(LIST), 4)
    print("k-th largest = %d" % v)
    print("List now: %s" % LIST)