题目描述:数据流中的第K大元素
解法一:排序
这个题目本身并不难解决,最容易想到的方法就是每次数据流中加入一个新元素后,就对整个数组从大到小排序,然后返回排序好的第 K 个数字即为数组中第 K 大的元素。
class KthLargest(object):
def __init__(self, k, nums):
"""
:type k: int
:type nums: List[int]
"""
self.k = k
self.nums = nums
self.nums.sort(reverse = True) # 从大到小排序
if len(self.nums) >= k:
self.kLargest = self.nums[k-1] # 如果len(nums)>=k,那第 k大元素就是nums[k-1]
else:
self.kLargest = None # 如果len(nums) <k,那目前还不存在第 k大的元素
def add(self, val):
"""
:type val: int
:rtype: int
"""
if val > self.kLargest: # 新元素只有大于当前数据流中的第 k大的元素才可能跻身前 k个数,否则第 k大的元素仍然不变,val永远不可能跻身前 k个数,无需加入排序数组。
self.nums.append(val)
self.nums.sort(reverse = True) # 重新排序
self.kLargest = self.nums[self.k-1] # 产生新的第 k大的元素
return self.kLargest
可想而知,这样做的时间复杂度肯定还是比较高的,因为每插入一个大于当前第 k大元素的新数据时,就要对新的整个数组重新排序一次。
改进版本:只维护一个长度为 k的数组,每次不需要再对整个数组进行排序,而是只对前 k个数重新进行排序。
class KthLargest(object):
def __init__(self, k, nums):
"""
:type k: int
:type nums: List[int]
"""
self.k = k
self.nums = nums
if len(nums) >= k:
self.nums.sort(reverse = True)
self.kLargest = self.nums[k-1]
def add(self, val):
"""
:type val: int
:rtype: int
"""
if len(self.nums) < self.k: #不够 k个元素时,先插入再排序
self.nums.append(val)
self.nums.sort(reverse = True)
self.kLargest = self.nums[self.k-1]
else:
if val > self.kLargest:
self.nums[self.k -1] = val # 新元素大于当前数据流中的第 k大的元素时,替换掉第 k大元素为val,保持数组中只有前 k个元素
self.nums.sort(reverse = True) # 重新对前 k个数排序
self.kLargest = self.nums[self.k-1]
return self.kLargest
解法二:借助小顶堆
其实,在前面的改进版本中,每次都只保留前 k个数,并只对这前 k大的数据进行维护,这其实可以借助堆这种数据结构来实现。具体来说,我们可以构建一个大小为 k 的小顶堆,保留整个数组中从大到小排序后的前 k个数,那么堆顶元素就是第 k 大的元素。每接收一个新数据,我们就拿它与当前堆顶元素比较,只要大于堆顶元素才可能加入堆,替换掉当前的堆顶元素,然后调整堆,重新找出新的第 k 大元素。
在 Python 中,内置的标准库提供了实现“二叉堆”数据结构的模块 heapq,其具体用法可以参考官方文档。heapq 其实是对一个list做原地的处理,第一个元素就是最小的,直接返回它就是小顶堆的堆顶元素。
class KthLargest(object):
import heapq
def __init__(self, k, nums):
"""
:type k: int
:type nums: List[int]
"""
# 题目中说明 len(nums) >= k-1,因此加入一个无穷小使得 len(nums) >= k,以便构造一个大小为 k 的小顶堆
# heapq.nlargest 返回数组前 k大个数(从大到小排好序的 K个数),返回类型为 list
self.topK_nums = heapq.nlargest(k, nums+[float('-inf')])
# 对前K大个数原地堆化,构建成一个小顶堆
heapq.heapify(self.topK_nums)
def add(self, val):
"""
:type val: int
:rtype: int
"""
# 每次插入新数据之后,就丢弃小于第 k 大元素的数据,相当于一直维持着堆的大小为 k
heapq.heappushpop(self.topK_nums, val)
# 返回topK个数中的最小元素即为第K大元素,也就是这K个数构成的小顶堆的堆顶元素
return self.topK_nums[0]
不调包,手动实现堆的写法
class KthLargest(object):
def __init__(self, k, nums):
"""
:type k: int
:type nums: List[int]
"""
# 题目中说明 len(nums) >= k-1,因此加入一个无穷小使得 len(nums) >= k,以便构造一个大小为 k 的堆
self.topK = sorted(nums+[float('-inf')], reverse=True)[:k] # 获取前K大个数
self.build_heap(self.topK) # 对前K大个数原地堆化,构建成一个小顶堆
def build_heap(self, nums):
# 堆是一棵完全二叉树,对于完全二叉树来说,下标从 n/2+1 到 n 的节点都是叶子节点
# 由于叶子节点往下堆化也只能和自己比较了
# 从倒数第二层第一个非叶子节点起,从后往前处理数组,并且每个数据都是从上往下堆化
for i in range(len(nums)//2, -1, -1):
self.heapify(nums, i)
def heapify(self, nums, idx): # 小顶堆的堆化过程
while True:
min_loc = idx
left = 2*idx + 1 # 针对下标起始为0的计算公式
right = 2*idx + 2
if left < len(nums) and nums[left] < nums[idx]: # 与左子节点比较
min_loc = left
if right < len(nums) and nums[right] < nums[min_loc]: # 与右子节点比较
min_loc = right
# 当前节点满足了父子节点之间的大小关系,确定了其最终所在位置,跳出循环
if min_loc == idx: break
# 否则,如果不满足子节点小于等于父节点的大小关系,我们就互换两个节点
tmp = nums[idx]
nums[idx] = nums[min_loc]
nums[min_loc] = tmp
# 当发生互换后,更新min_loc为当前节点与其子节点交换后的所在位置,循环继续比较子节点与子子节点的大小关系(往下堆化)
idx = min_loc
def add(self, val):
"""
:type val: int
:rtype: int
"""
if self.topK[0] < val: # 若新插入的值大于堆顶元素,则替换堆顶元素并调整堆
self.topK[0] = val
self.heapify(self.topK, 0)
return self.topK[0] # topK个数构成的小顶堆的堆顶元素即为第K大元素
另一种写法
class KthLargest(object):
def __init__(self, k, nums):
"""
:type k: int
:type nums: List[int]
"""
self.k = k
self.heap = nums
heapq.heapify(self.heap) # heapq.heapify(x)可将一个 list 原地构建成小顶堆,这里传入的 heap其实就是个list
while len(self.heap) > self.k: # 将堆中元素减小到只有前 k大个数
heapq.heappop(self.heap) # 通过不断删除堆顶元素,并重新调整堆
def add(self, val):
"""
:type val: int
:rtype: int
"""
if len(self.heap) < self.k: # 堆中不够 k个元素,则直接添加 val进去
heapq.heappush(self.heap, val) # 新增元素到堆中
else:
if self.heap[0] < val: # 若新的值大于堆顶元素
heapq.heapreplace(self.heap, val) # 则删除堆顶元素,并重新堆化前 K大个数构成的小顶堆
return self.heap[0] # 最终返回堆顶元素,即为第K大元素
数组中的第K个最大元素
解法一:直接从大到小排序,取第k个元素
class Solution(object):
def findKthLargest(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
return sorted(nums, reverse=True)[k-1]
# 或:return sorted(nums)[-k]
时间复杂度:采用快排即为 O(n*logn),n 为数组长度
空间复杂度:O(1)
解法二:维护一个大小为k的小顶堆
class Solution(object):
def findKthLargest(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
# 先取数组中前k个数构建一个小顶堆
kLargest = nums[:k]
self.build_heap(kLargest)
# 从第k+1个数起,遍历数组中剩下的元素
for i in range(k, len(nums)):
if kLargest[0] < nums[i]: # 如果遍历到的元素大于当前堆顶元素,则替换堆顶元素并重新调整小顶堆
kLargest[0] = nums[i]
self.heapify(kLargest, 0)
return kLargest[0] # 遍历完返回topK个数构成的小顶堆的堆顶元素即为第K大元素
def build_heap(self, nums): # 建堆
for i in range(len(nums)//2, -1, -1):
self.heapify(nums, i)
def heapify(self, nums, idx): # 小顶堆的堆化过程
while True:
min_pos = idx
left = idx*2 + 1
right = idx*2 + 2
if left<len(nums) and nums[left] < nums[idx]:
min_pos = left
if right<len(nums) and nums[right] < nums[min_pos]:
min_pos = right
if min_pos == idx: break
nums[idx], nums[min_pos] = nums[min_pos], nums[idx] # 互换
idx = min_pos
调用Python 的 heapq 库的写法,可以精简到一行代码(果然是人生苦短,我用python,手动狗头…)
class Solution(object):
def findKthLargest(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
import heapq
return heapq.nlargest(k,nums)[-1]
时间复杂度:往大小为 k 的堆中添加元素的时间复杂度为 O(logk),重复该操作 n 次,故总时间复杂度为 O(n*logk)
空间复杂度:维护一个存放前 k 个元素的堆需要 O(k)
解法三:快速选择算法
本方法思路大致上与快速排序相同。简便起见,注意到第 k 个最大元素也就是第 N - k 个最小元素,因此可以用求第 N - k 小的元素的思路来解决本问题。
- 随机选择一个枢纽
- 沿着数组移动,将每个元素与枢轴进行比较,并将小于枢轴的所有元素移动到枢轴的左侧,所有大于或等于的元素都放在其右侧,最终在输出的数组中,枢轴落在了其合适的位置 pos 上。
- 比较 pos 和 N - k 以决定在哪边继续递归处理
class Solution(object):
def findKthLargest(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
left = 0
right = len(nums)-1
k_smallest = len(nums)-k # 第n-k小的元素,即为第k大元素
return self.quick_select(nums, left, right, k_smallest)
def partition(self, nums, left, right, pivot_index):
pivot = nums[pivot_index]
# 1.先把枢纽元素换到数组末尾上存放
nums[pivot_index], nums[right] = nums[right], nums[pivot_index]
# 2.通过游标sorted_index把nums分成两部分。
# nums[left:sorted_index]中的元素都是小于pivot的,可以理解为已处理区间,剩下的nums[sorted_index:right]为未处理区间
# 每次从未处理区间中取一个元素与pivot对比,如果小于pivot则将其加入到已处理区间的尾部,也就是sorted_index位置上。
sorted_index = left
for i in range(left, right):
if nums[i] < pivot:
nums[i], nums[sorted_index] = nums[sorted_index], nums[i]
sorted_index += 1
# 3.将pivot放在其最终合适的位置上
nums[right], nums[sorted_index] = nums[sorted_index], nums[right]
return sorted_index # 返回最终的pivot_index
def quick_select(self, nums, left, right, target):
if left == right: # 递归终止条件
return nums[left]
pivot_index = random.randint(left, right) # 每次随机选择一个位置作为分区的枢纽
pivot_index = self.partition(nums, left, right, pivot_index) # 调用分区算法,确定排序后pivot_index合适的位置
if target == pivot_index: # 当前枢纽所在位置正好为待查找的位置
return nums[pivot_index]
elif target < pivot_index: # 如果待查找的位置小于当前枢纽所在位置,去左半边递归查找
return self.quick_select(nums, left, pivot_index-1, target)
else: # 如果待查找的位置大于当前枢纽所在位置,去右半边递归查找
return self.quick_select(nums, pivot_index+1, right, target)
也可以对数组按倒序进行分区,然后直接查找第k大元素
class Solution(object):
def findKthLargest(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
left = 0
right = len(nums)-1
return self.quick_select(nums, left, right, k-1)
def partition(self, nums, left, right, pivot_index):
pivot = nums[pivot_index]
nums[pivot_index], nums[right] = nums[right], nums[pivot_index]
# 每次从未处理区间中取一个元素与pivot对比,如果大于pivot则将其放到pivot_index左边。
sorted_index = left
for i in range(left, right):
if nums[i] > pivot:
nums[i], nums[sorted_index] = nums[sorted_index], nums[i]
sorted_index += 1
nums[right], nums[sorted_index] = nums[sorted_index], nums[right]
return sorted_index
def quick_select(self, nums, left, right, target):
if left == right: # 递归终止条件
return nums[left]
pivot_index = random.randint(left, right) # 每次随机选择一个位置作为分区的枢纽
pivot_index = self.partition(nums, left, right, pivot_index) # 调用分区算法,确定排序后pivot_index合适的位置
if target == pivot_index: # 当前枢纽所在位置正好为待查找的位置
return nums[pivot_index]
elif target < pivot_index: # 如果待查找的位置小于当前枢纽所在位置,去右半边递归查找
return self.quick_select(nums, left, pivot_index-1, target)
else: # 如果待查找的位置大于当前枢纽所在位置,去左半边递归查找
return self.quick_select(nums, pivot_index+1, right, target)
时间复杂度:平均情况下为O(N),最坏情况下为O(N2)
空间复杂度:O(1)
最接近原点的 K 个点
解法一:直接排序
按每个点与原点的距离从小到大排序,然后取前 k 个点即为距离原点最近的 k 个点。
class Solution(object):
def kClosest(self, points, K):
"""
:type points: List[List[int]]
:type K: int
:rtype: List[List[int]]
"""
distance_dict = {}
for point in points:
distance = sqrt(point[0]**2+point[1]**2)
distance_dict[(point[0],point[1])] = distance
sorted_dict = sorted(distance_dict.items(), key =lambda x:x[1]) # 按字典的value排序,输出类型为list
res = []
for i in range(K): # 取排序好的前K个点即为距离原点最近的K个点
res.append(sorted_dict[i][0]) # sorted_dict中每一项为tuple,元组的第一个元素即为坐标点
return res
另一种写法,可以利用python语言的匿名函数特性,通过自定义一个排序的键值比较函数来完成排序
class Solution(object):
def kClosest(self, points, K):
"""
:type points: List[List[int]]
:type K: int
:rtype: List[List[int]]
"""
sorted_points = sorted(points, key =lambda x: sqrt(x[0]**2+x[1]**2))
return sorted_points[:K] # 前K个距离原点最近的点
时间复杂度:O(N*logN)
空间复杂度:O(1)
解法二:维护一个大小为k的大顶堆
class Solution(object):
def kClosest(self, points, K):
"""
:type points: List[List[int]]
:type K: int
:rtype: List[List[int]]
"""
import heapq
# 取出K个距离原点最近的点
k_smallest = heapq.nsmallest(K, points, key=lambda p: sqrt(p[0]**2+p[1]**2))
return k_smallest
不调用库,手动实现堆排序的写法
class Solution(object):
def kClosest(self, points, K):
"""
:type points: List[List[int]]
:type K: int
:rtype: List[List[int]]
"""
# 先取数组前K个点构建一个大顶堆(根据元素与原点的距离比较大小)
K_smallest = points[:K]
self.build_heap(K_smallest)
# 从第K+1个元素开始遍历剩余的点,如果发现有比当前堆顶元素距离原点更近的点,则替换掉堆顶元素
for i in range(K, len(points)):
if self.calDist(K_smallest[0]) > self.calDist(points[i]):
K_smallest[0] = points[i]
self.heapify(K_smallest, 0)
return K_smallest
def build_heap(self, nums): # 建堆
for i in range(len(nums)//2, -1, -1): # 从后往前遍历数组中的元素,并且每个元素都是从上往下堆化
self.heapify(nums, i)
def heapify(self, nums, idx): # 大顶堆的堆化过程
while True:
max_loc = idx
left = idx*2+1
right = idx*2+2
if left<len(nums) and self.calDist(nums[left]) > self.calDist(nums[idx]):
max_loc = left
if right<len(nums) and self.calDist(nums[right]) > self.calDist(nums[max_loc]):
max_loc = right
if max_loc == idx: break
nums[idx], nums[max_loc] = nums[max_loc], nums[idx]
idx = max_loc
def calDist(self, point):
return sqrt(point[0]**2+point[1]**2)
时间复杂度:O(N*logK)
空间复杂度:O(K)
解法三:快速选择法
class Solution(object):
def kClosest(self, points, K):
"""
:type points: List[List[int]]
:type K: int
:rtype: List[List[int]]
"""
def partition(points, left, right, pivot_idx):
pivot = self.calDist(points[pivot_idx])
# 先把pivot的值与最右边的值交换,存放到最右边
points[pivot_idx], points[right] = points[right], points[pivot_idx]
# 把所有小于pivot的元素都放到左半边
sorted_idx = left
for i in range(left, right):
if self.calDist(points[i]) < pivot:
points[i], points[sorted_idx] = points[sorted_idx], points[i]
sorted_idx += 1
# 把pivot放到其最终合适的位置上
points[sorted_idx], points[right] = points[right], points[sorted_idx]
return sorted_idx
def quick_select(points, left, right, K):
if left == right: return # 递归终止条件,未排序区间只剩下一个元素
# 1.随机选择一个位置作为分区pivot
pivot_idx = random.randint(left, right)
# 2.调用partition函数,确定pivot最终的位置
pivot_idx = partition(points, left, right, pivot_idx)
# 3.判断最小的K个值落在哪个区间中
if K < pivot_idx-left+1: # 最小的K个值均落在pivot左侧,递归对左半边排序
quick_select(points, left, pivot_idx-1, K)
elif K > pivot_idx-left+1: # 已经得到(pivot_idx-left+1)个最小值,但不足K个,递归到右半边寻找剩余的K-(pivot_idx-left+1)个最小值
quick_select(points, pivot_idx+1, right, K-(pivot_idx-left+1))
else: # 此时pivot==K,左半边的所有值即为最小的K个值
return
quick_select(points, 0, len(points)-1, K) # 调用快速选择算法,找出前K个最小值
return points[:K]
def calDist(self, point): # 计算与原点的距离
return sqrt(point[0]**2+point[1]**2)
时间复杂度:平均情况下为O(N),最坏情况下为O(N2)
空间复杂度:O(1)
最后,如果大家有更好的Python解法,欢迎在评论区分享下您的解法,一起进步,感谢^ o ^~
参考:
https://leetcode-cn.com/problems/kth-largest-element-in-a-stream/solution/python-3xing-er-fen-otlogn-4-xing-dui-otlogk-by-qq/
https://leetcode-cn.com/problems/kth-largest-element-in-an-array/solution/shu-zu-zhong-de-di-kge-zui-da-yuan-su-by-leetcode/