了解递归函数-快速选择-在线性时间内查找中位数

Question

I am trying to implement this algorithm (from this site: https://sarielhp.org/research/CG/applets/linear_prog/median.html ).

FindKMedian( A, K ) // Return the number in A which is the K-th in its size.

1. Pick randomly a number a from A = {a1, ..., an}.
2. Partition the n numbers into two sets:
   • S - all the numbers smaller than a
   • B - all the numbers bigger than a
3. If |S| = K-1 then a is the required K-median. Return a
4. If |S| < K-1 then the K-median lies somewhere in B. Call recursively to FindKMedian( B, K - |S| - 1 )
5. Else, call recursively to FindKMedian( S, K ).

After @mikake answer I get an error calling the function with the parameters at the end of the code.

import random


def quick_select(A, k, s, d):
    r = random.choice(range(s, d))
    pivot = partition(A, r, s, d)
    if pivot == k:
        return A[pivot]
    elif k < pivot:
        return quick_select(A, k, s, pivot-1)
    else:
        return quick_select(A, k, pivot + 1, d)


def partition(A, r, s, d):
    j = s-1
    assert s <= r
    assert r <= d
    temp = A[d]
    A[d] = A[r]
    A[r] = temp
    for i in range(s, d):
        if A[i] <= A[d]:
            j = j+1
            temp = A[j]
            A[j] = A[i]
            A[i] = temp
    j = j+1
    temp = A[j]
    A[j] = A[d]
    A[d] = temp
    return j


random.seed(0)
A = [4, 7, 7, 2, 2, 0, 9, 8, 1, 8]
print(quick_select(A, 5, 0, 9))

I would expect the number 7 to come out from the return of quickselect (so quick_select(A, 5, 0, 9) means "find A[5] once the subarray A[0,...,5] is sorted or once A[5,...,9] is sorted "). I probably didn't get what the semantic of this code should be.

Thank you

Answer 1

You forgot to add the return statement in the "else" branches:

def quick_select(A, k, s, d):
    r = random.choice(range(s, d))
    pivot = partition(A, r, s, d)
    if pivot == k:
        return A[pivot]
    elif k < pivot:
        return quick_select(A, k, s, pivot-1)
    else:
        return quick_select(A, k, pivot + 1, d)

Answer 2

I think the only error I made was not considering the case when the array has length 1. So the correct code of the function "quick_select" should be

def quick_select(A, k, s, d):
    if s == d:
        return A[k]
    r = random.choice(range(s, d))
    pivot = partition(A, r, s, d)
    if pivot == k:
        return A[pivot]
    elif k < pivot:
        return quick_select(A, k, s, pivot-1)
    else:
        return quick_select(A, k, pivot + 1, d)