leetcode__有序矩阵中第K小的元素__python

解题思路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]

你可能感兴趣的:(leetcode,python)