leetcode有序矩阵中第K小的元素python

最新推荐文章于 2025-05-19 05:00:00 发布

原创最新推荐文章于 2025-05-19 05:00:00 发布 · 367 阅读

0 ·

CC 4.0 BY-SA版权

leetcode 同时被 2 个专栏收录

83 篇文章

订阅专栏

python

66 篇文章

订阅专栏

本文介绍了两种求解矩阵中第K小元素的方法。一种是将所有元素整合成一个列表并排序；另一种是使用最大堆的数据结构进行高效查找。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

解题思路1：

直接将所有的元素合成一个列表，排序后取出第k小的元素即可。
具体代码如下：

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        my_list = []
        for tmp in matrix:
            my_list += tmp
        my_list = sorted(my_list)
        return my_list[k-1]

解题思路2：

利用最大堆的思想解题即可。
具体代码如下:

class max_heap:
    def __init__(self):
        self.data = []
        
    def get_parent(self, idx):
        return (idx-1)//2
    
    def push(self, num):
        self.data.append(num)
        ptr = len(self.data) - 1
        parent_id = self.get_parent(ptr)
        while self.data[parent_id] < self.data[ptr]:
            tmp = self.data[ptr]
            self.data[ptr] = self.data[parent_id]
            self.data[parent_id] = tmp
            if parent_id == 0:
                break
            else:
                ptr = parent_id
                parent_id = self.get_parent(ptr)
    
    def pop(self):
        self.data = [self.data[-1]] + self.data[1:-1]
        #has left
        ptr = 0
        while 2*ptr + 1 < len(self.data):
            #has right
            if 2*ptr + 2 < len(self.data):
                if self.data[2*ptr + 1] > self.data[2*ptr + 2]:
                    max_idx = 2*ptr + 1
                else:
                    max_idx = 2*ptr + 2
                if self.data[max_idx] > self.data[ptr]:
                    tmp = self.data[ptr]
                    self.data[ptr] = self.data[max_idx]
                    self.data[max_idx] = tmp
                    ptr = max_idx
                else:
                    break
            else:
                if self.data[2*ptr + 1] > self.data[ptr]:
                    tmp = self.data[ptr]
                    self.data[ptr] = self.data[2*ptr + 1]
                    self.data[2*ptr + 1] = tmp
                    ptr = 2*ptr + 1 
                else:
                    break        


class Solution:
    def kthSmallest(self, matrix, k):
        my_list = max_heap()
        for sample_list in matrix:
            for num in sample_list:
                if len(my_list.data) < k:
                    my_list.push(num)
                else:
                    if num < my_list.data[0]:
                        my_list.pop()
                        my_list.push(num)
        return my_list.data[0]