数据结构与算法：寻找数组第K大元素及其python实现_k largest 算法

一、简单排序

最简单的想法就是，将原数组降序排序，再取其第k-1个元素。此种算法的时间复杂度为$O(N\log N)$。

二、多次查找最大元素

每次找到数组的最大元素，复杂度$O(N)$，然后再对剩下的元素找到最大，重复$K$次，总时间复杂度为$O(K\ast N)$；

三、快排思想

若对目标数组使用快速排序（降序），快排在每次迭代时，选取一个元素作为枢纽，一次迭代完成后，枢纽元素不小于其左侧元素且不大于其右侧元素。因此，在每次迭代完成时，检查枢纽元素的位置，即

• 当枢纽元素位置刚好等于数组中第K大元素对应的位置（K-1），则找到目标元素；
• 当枢纽位置小于K-1，则在枢纽右侧的子序列中查找目标元素；
• 当枢纽位置大于K-1，则在枢纽左侧的子序列中查找目标元素；

算法的时间复杂度不是$O(N\log N)$，因此第一次迭代平均交换N次，第二次迭代平均交换N/2次，依次类推，最后一次迭代交换1次，因此平均交换次数为
$$N+\frac{N}{2}+\frac{N}{4}+\cdots <2N$$

故时间复杂度为$O(N)$。

def __partition(arr, start, end):
    """部分逆序"""
    pivot = arr[start]
    while start < end:
        while start < end and pivot >= arr[end]:
            end -= 1
        if start < end:
            arr[start] = arr[end]
            start += 1
        while start < end and pivot <= arr[start]:
            start += 1
        if start < end:
            arr[end] = arr[start]
            end -= 1
    arr[start] = pivot
    return start


def find_k_largest(arr, k):
    """查找数组中第k大的数"""
    start, end = 0, len(arr) - 1
    assert 0 < k <= end + 1

    while True:
        pos = __partition(arr, start, end)
        if pos + 1 == k:
            return arr[pos]
        elif pos < k:
            start = pos + 1
        else:
            end = pos - 1


if __name__ == '__main__':
    print(find_k_largest([1, 1, 2, 3, 20, 4, 10, 11], 2))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

四、计数思想

若目标数组为整型数组且取值范围较小，则可通过统计数组中各数据的频数，找到第K大元素。

首先统计目标数组的最大最小值为maxV和minV，分配长度为maxV-minV+1的辅助数组用于存储各元素的频数。遍历目标数组，统计各元素频数，然后将各元素频数从后向前依次累加，当累加值刚好大于等于K时，则当前元素即为目标第K大元素。

时间复杂度为$O(N+maxV-minV+1)$，空间复杂度为$O(maxV-minV+1)$。

def min_max(arr):
    """查找数组中最小、最大值"""

    return min, max


def find_k_largest(arr, k):
    """查找数组中第k大的数"""
    min = max = arr[0]
    for num in arr[1:]:
        if num > max:
            max = num
        elif num < min:
            min = num

    count = [0] * (max - min + 1)
    for num in arr:
        count[num - min] += 1

    total = 0
    for i in range(max - min, -1, -1):
        total += count[i]
        if total >= k:
            return i + min


if __name__ == '__main__':
    print(find_k_largest([1, 1, 2, 3, 20, 4, 10, 11], 2))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

五、最小堆思想

初始时构建K个元素的最小堆，再遍历目标数组中的剩余元素，若元素大于堆顶元素，则将堆顶元素覆盖，并重建最小堆。最后所得最小堆的堆顶元素，即为第$K$大目标元素。

初始K个元素的最小堆的时间复杂度为$O(K)$，遍历目标数组中的剩余元素并重建最小堆的复杂度为$O((N-K)\log K)$，故总时间复杂度为$O(K+(N-K)\log K)$。

def min_heapfiy(arr, start, end):
    """小堆化"""
    son = 2 * start + 1
    while son <= end:
        if son < end and arr[son] > arr[son + 1]:
            son += 1
        if arr[start] <= arr[son]:
            return
        arr[start], arr[son] = arr[son], arr[start]
        start, son = son, son * 2 + 1


def find_k_largest(arr, k):
    """查找数组中第k大元素"""
    assert 0 < k <= len(arr)

    heap = [num for num in arr[:k]]
    for start in range((k >> 1) - 1, -1, -1):
        min_heapfiy(heap, start, k - 1)

    for num in arr[k:]:
        if num > heap[0]:
            heap[0] = num
            min_heapfiy(heap, 0, k - 1)

    return heap[0]


if __name__ == '__main__':
    print(find_k_largest([0, 5, 2, 3, 4, 100, 30, 105, 60, 7, 10, 7, 1], 2))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30