algorithm-in-python/algorithm/sort/select.py

''' mbinary
#########################################################################
# File : select.py
# Author: mbinary
# Mail: zhuheqin1@gmail.com
# Blog: https://mbinary.coding.me
# Github: https://github.com/mbinary
# Created Time: 2018-07-06  17:13
# Description:
#########################################################################
'''

from random import randint
def select(lst,i):
    lst = lst.copy()
    def partition(a,b):
        pivot = lst[a]
        while a<b:
            while a<b and lst[b]>pivot: b-=1
            if a<b:
                lst[a] = lst[b]
                a+=1
            while a<b and lst[a]<pivot: a+=1
            if a<b:
                lst[b] = lst[a]
                b-=1
        lst[a]= pivot
        return a

    def _select(a,b):
        if a>=b: return lst[a]
        # randomized select
        n = randint(a,b)
        lst[a],lst[n] = lst[n],lst[a]
        pos = partition(a,b)
        if pos>i:
            return _select(a,pos-1)
        elif pos<i:
            return _select(pos+1,b)
        else:return lst[pos]
    return _select(0,len(lst)-1)


if __name__ =='__main__':
    lst = [randint(0,1000) for i in range(100)]
    st = sorted(lst)
    for i in range(10):
        n = randint(0,99)
        print('select {}th: \nexpect: {}\ngot: {}'.format(n,st[n],select(lst,n)))