randomized selection algorithm

selection.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# pylint: disable=C0111
"""
Randomized Selection Algorithm:
    O(n) for query                (expected; O(n^2) for worstcase)
    O(1) for insertion
    O(n) for deletion             (expected; O(n^2) for worstcase)

    this algo is good if there are less queries and more insertions
    on an average this is not good for online type problems.

Two Priority Queue Method:
    O(log(n)) for query
    O(log(n)) for insertion
    O(log(n)) for deletion

    For online type problems you can use two priority queue method.
    which takes O(log(n)) for insertion and O(log(n)) for query and btw
    it also supports the deletion of the queried element in O(log(n))
"""
from random import shuffle


def _partition(arr, lo, hi, pivot):
    p = arr[pivot]
    arr[hi - 1], arr[pivot] = arr[pivot], arr[hi - 1]
    i = lo - 1
    for j in xrange(lo, hi):
        if arr[j] <= p:
            i += 1
            arr[i], arr[j] = arr[j], arr[i]
    return i


def select(arr, lo, hi, spos):
    assert lo <= spos < hi
    shuffle(arr)  # here's your randomization.
    while True:
        pos = _partition(arr, lo, hi, lo)
        if pos == spos:
            return arr[pos]
        elif pos < spos:
            lo = pos + 1
        else:
            hi = pos