线性时间求取第 K 大数

Posted 2020-11-24 seniusen

求 Top K 的算法主要有基于快速排序的和基于堆的这两种，它们的时间复杂度都为 (O(nlogK))。借助于分治思想，以及快速排序的区间划分，我们可以做到 (O(n)) 时间复杂度。具体算法思路如下：

• 第 1 步，我们将原数据 5 个一组划分为若干个组，最后余下的不足 5 个的额外作为一组，总组数为 (g=lceil{n/5} ceil)；
• 第 2，3 步， 对每一个组内的 5 个元素利用插入排序算法进行排序，然后将每个组的中位数依次放到数据的前面，最后 (A[0, g-1]) 保存的便是 (g) 个组各自的中位数；
• 第 4 步，递归调用求出 (A[0, g-1]) 的中位数，也即是元素 (A[g/2])，不知是否可以直接得到；
• 第 5 步，以上一步得到的中位数对数据划分区间，左边小于中位数，右边大于中位数，中位数位于第 (k) 个位置；
• 第 6 步，如果我们要找的正好是第 (k) 大数，那么刚刚得到的中位数即为所求；
• 第 7 步，如果 (i<k) 我们要找的第 (i) 大数位于区间 (A[0, k-1])，在左半边递归求取第 (i) 大的数；
• 第 8 步，如果 (i>k) 我们要找的第 (i) 大数位于区间 (A[k+1:-1])，在右半边递归求取第 (i-k) 大的数。

def insert_sort(data, left, right):

    # 对列表 data[left, right] 进行插入排序

    for i in range(left+1, right+1):
        num = data[i]
        for j in range(i, left, -1):
            if num < data[j-1]:
                data[j] = data[j-1]
            else:
                break
        data[j] = num


def swap(data, i, j):

    # 交换列表 data 位于 i，j 的元素

    temp = data[i]
    data[i] = data[j]
    data[j] = temp


def partition(data, x):

    # 按照列表 data 第 x 个元素分区

    n = len(data)
    i = 0
    pivot = data[x]
    swap(data, x, n-1)
    for j in range(0, n):
        if data[j] < pivot:
            swap(data, i, j)
            i += 1

    data[j] = data[i]
    data[i] = pivot

    return i


def select_ith_min(data, i):

    # 选取列表 data 中第 i 小的元素

    n = len(data)

    # 不要忘了递归终止
    if n == 1:
        return data[0]

    if n % 5 == 0:
        group = n // 5
    else:
        group = n // 5 + 1
    for j in range(0, group):
        start = j * 5
        end = start + 4
        if end > n - 1:
            end = n - 1
        insert_sort(data, start, end)
        mid = (end - start + 1) // 2 + start
        swap(data, j, mid)

    pivot = data[group // 2]
    k = partition(data, group // 2)

    if k == i-1:
        return pivot
    elif k > i-1:
        return select_ith_max(data[0:k], i)
    else:
        return select_ith_max(data[k:], i-k)


a = [i for i in range(100, 0, -1)]
print(a)
for i in range(1, 101):
    num = select_ith_max(a, i)
    print(num)

获取更多精彩，请关注「seniusen」!