"""
Selection algorithms

python implementaiton of quickselect and heapselect

quickselect does not guarantee the selection to be sorted and has a O(n)
complexity, while heapselect sorts the selection in O(n*log(k))
"""

import random
import heapq

def quickselect(l, k):
    pivot = random.choice(l)
    pivots = [x for x in l if x == pivot]
    left = [x for x in l if x < pivot]
    right = [x for x in l if x > pivot]

    if len(right) > k:
        return quickselect(right, k)
    if len(right) + len(pivots) >= k:
        return right + pivots[:(k - len(right))]
    return quickselect(left, k - len(right) - len(pivots)) + pivots + right


def heapselect(l, k):
    return heapq.nlargest(k, l)