欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页  >  IT编程

python查找第k小元素代码分享

程序员文章站 2022-06-06 19:16:05
复制代码 代码如下:# -*- coding: utf-8 -*- from random import randintfrom math import ceil, fl...

复制代码 代码如下:

# -*- coding: utf-8 -*-

from random import randint
from math import ceil, floor

def _partition(a, l, r, i):
    """以a[i]为主元划分数组a[l..r],使得:
    a[l..m-1] <= a[m] < a[m+1..r]
    """
    a[i], a[r] = a[r], a[i] # i交换到末位r,作为主元
    pivot = a[r] # 主元
    m = l # 索引标记
    for n in xrange(l, r): # l..r-1
        if a[n] <= pivot:
            a[m], a[n] = a[n], a[m] # 交换
            m += 1 # 后移
    a[m], a[r] = a[r], a[m] # 主元到m位
    return m

def _rand(a, l, r):
    """随机划分主元"""
    return randint(l, r) # a[l..r]随机取一个

def _select(a, l, r, k, pivot_selector = _rand):
    """利用快排,得a[l..r]中第k小的数,k in [l+1,r+1]:

    其尾递归方式,伪码如下:
    select(a, l, r, k)
    1  while true:
    2    i ← ? // 划分主元位置
    3    m ← partition(a, l, r, i) // 数组划分
    4    n ← m - l + 1 // a[l..m]元素个数
    5    if k = n // 检查a[m]是否是第k小的元素
    6      then return a[m]
    7    elseif k < n // 左划分区
    8      r = m - 1
    9    else // 右划分区
    10     k = k - n
    11     l = m + 1

    args:
        pivot_selector(function): 主元选取方法,默认随机方式
    """
    if not a:
        return none
    if l == r:
        return a[l]
    while true:
        i = pivot_selector(a, l, r)
        m = _partition(a, l, r, i)
        n = m - l + 1
        if k == n:
            return a[m]
        elif k < n:
            r = m - 1
        else:
            k = k - n
            l = m + 1

def rand_select(a, k):
    """默认随机划分主元方式,k in [1, len(a)]
    e[t(n)] = o(n)
    """
    return _select(a, 0, len(a) - 1, k);


def _median(a, l, r):
    """对a[l..r]插入排序(原地)后选取其中位数位置"""
    for j in xrange(l, r + 1):
        k = a[j]
        i = j
        while i > l and a[i-1] > k:
            a[i] = a[i-1]
            i -= 1
        a[i] = k
    return l + int((r - l) * 0.5) # 下中位数

def _medianofmedians(a, l, r):
    """中位数的中位数方式:
    1. 划分为floor(n/5)个5元组,剩下(n%5)组成最后一组。
    2. 找出ceil(n/5)个组各自的中位数。先对每组插入排序,再从中选出中位数。
    3. 对第2步中找出的ceil(n/5)个中位数重复上述操作,直到仅有一个中位数。
    """
    if l == r:
        return l
    n = r - l + 1 # 元素个数
    m = int(ceil(n / 5.0)) # 划分组数,每组5个元素
    for i in xrange(m):
        # 每组起始位和结束位
        sub_l = l + i * 5
        sub_r = sub_l + 4
        if sub_r > r:
            sub_r = r
        # 对每组元素插入排序后,选取中位数
        sub_m = _median(a, sub_l, sub_r) # 中位数索引
        # 交换中位数到前几位
        j = l + i
        a[j], a[sub_m] = a[sub_m], a[j]
    return _medianofmedians(a, l, l + m - 1) # 中位数的中位数

def bfprt_select(a, k):
    """中位数的中位数方式(bfprt算法)
    t(n) = o(n)
    """
    return _select(a, 0, len(a) - 1, k, _medianofmedians);


def _median3(a, l, r):
    """三数中位数方式,取l,r,(l+r)/2三数中位数"""
    c = (l + r) / 2
    keys = [l, c, r]
    i = _median(keys, 0, 2)
    return keys[i]

def median_select(a, k):
    """三数中位数方式,以消除最坏情况"""
    return _select(a, 0, len(a) - 1, k, _median3);


if __name__ == '__main__':
    import random, time
    from copy import copy

    print('preparing data...')
    n = 1000000
    nums = range(n)
    random.shuffle(nums)
    print('ready go!')

    def timeit(fnc, *args, **kargs):
        print('%s starts processing' % fnc.__name__)
        begtime = time.clock()
        retval = fnc(*args, **kargs)
        endtime = time.clock()
        print('%s takes time : %f' % (fnc.__name__, endtime - begtime))
        return retval

    test_methods = [rand_select, bfprt_select, median_select]
    k = random.randrange(n) + 1
    dashes = '---' * 10
    for test in test_methods:
        print(dashes)
        nums_new = copy(nums)
        result = timeit(test, nums_new, k)
        print('the %dth smallest element: %d' % (k, result))