Top K 问题

一步解决数组的Top K有点困难,不妨逐步解决,先从简单的开始,比如只要找到数组中的最大的一个数可以这样实现:

方法一:

1
2
3
4
5
6
7
8
9
10
def findLargest(arr, arr_size):
if not arr:
return -1

largest = -float('inf')
for i in range(arr_size):
if arr[i] > largest:
largest = arr[i]

return largest

增大难度,如果需要找到数组中最大的三个数,那又如何实现呢?

方法二:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def find3Largest(arr, arr_size):
if arr_size < 3:
print('无效的输入')
return -1

first = second = third = -float('inf')

for i in range(arr_size):
if arr[i] > first: # 找出数组中的最大值
third = second # 一旦数组中有新的最大值出现,那么原来*第二大*的就变成*第三大*
second = first # *第一大*的变成*第二大*
first = arr[i]

elif arr[i] > second:
third = second
second = arr[i]

elif arr[i] > third:
third = arr[i]

return first, second, third

找到最大的3个数也比较简单,那么问题来了,如果要找到最大的K个数,并且K比3要大,比如说10,仍旧用上面的办法就比较不现实。

方法三:冒泡排序变种

还记得冒泡排序吗?每一轮经过两两比较后会产生一个最大的数,如果我们控制外部循环的次数就能找到想要数组中相应最大的K个数,实现如下:

1
2
3
4
5
6
7
def findKLargest(arr, arr_size, K):

for i in range(arr_size, K+1, -1):
for j in range(i):
if arr[j] > arr[j+1]:
arr[j], arr[j+1] = arr[j+1], arr[j]
return arr[-K:]

缺点还是有的,时间复杂度太高了,最坏的情况下有$O(nk)$,所以当数组长度很长的时候肯定是不行的。

方法四:用堆排序(最大堆)

回忆堆排序的过程,堆排序在成功构建了最大堆之后,最后一步是从堆中取出最大的元素,然后再构建最大堆,再取出最大元素,所以我们也可以通过这种方法来实现从数组中找出K个最大元素的目标。构建最大堆的时间复杂度为$O(n)$,从堆中k次取出最大元素的复杂度为$O(klog(n)$。

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
36
37
38
39
40
41
42
43
44
# Python program for implementation of heap Sort 

# To heapify subtree rooted at index i.
# n is size of heap
# heapify()函数的作用是对索引i下的子树构建最大堆
def heapify(arr, n, i): # n是堆的大小,先假设i是根节点
largest = i # Initialize largest as root
l = 2 * i + 1 # left = 2*i + 1
r = 2 * i + 2 # right = 2*i + 2

# See if left child of root exists and is
# greater than root
if l < n and arr[i] < arr[l]:
largest = l

# See if right child of root exists and is
# greater than root
if r < n and arr[largest] < arr[r]:
largest = r

# Change root, if needed
if largest != i:
arr[i], arr[largest] = arr[largest], arr[i] # 如果发现左右子节点中有大于根节点的,则互换其位置

# Heapify the root.
heapify(arr, n, largest) # 继续对子树构建最大堆

# The main function to sort an array of given size
def heapSort(arr):
n = len(arr)

# Build a maxheap.
for i in range(n, -1, -1): # 从完全二叉树最低端的叶子节点往根节点遍历,构建最大堆,时间复杂度为O(n)
heapify(arr, n, i)

# One by one extract elements
for i in range(n-1, n-1-K, -1): # 数组中最大的元素永远在堆的最顶端,下标为0
arr[i], arr[0] = arr[0], arr[i] # 找出树中的最大元素,与最后一个node交换
heapify(arr, i, 0) # 重新构建最大堆,此时i是数组的长度,已经不包含最后一个最大的元素了

# Driver code to test above
arr = [ 12, 11, 13, 5, 6, 7]
heapSort(arr)
print ("Sorted array is", arr[-K:])
觉得还不错?赞助一下~
0%