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

代码如下:

# -*- 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