LeetCode 378 [Kth Smallest Sum i

原题

在一个排序矩阵中找从小到大的第 k 个整数。
排序矩阵的定义为：每一行递增，每一列也递增。

样例
给出 k = 4和一个排序矩阵：

[
  [1 ,5 ,7],
  [3 ,7 ,8],
  [4 ,8 ,9],
]

返回 5。

解题思路

• 首先把左上角的元素放入minHeap中，进入while循环，每次pop一个最小值，然后把该位置右边和下班的值+坐标放入minHeap中。k减一，当k等于0时，返回当前的值
• 同时要注意，另外开一个二维数组记录哪些元素已经被访问过，因为同一个元素可能在A的下面和B的右面

完整代码

import Queue

class Solution(object):
    def kthSmallest(self, matrix, k):
        """
        :type matrix: List[List[int]]
        :type k: int
        :rtype: int
        """
        exist = [[False for i in range(len(matrix[0]))] for j in range(len(matrix))]
        q = Queue.PriorityQueue()
        q.put((matrix[0][0], 0, 0))
        exist[0][0] = True
        while k > 0:
            cur, x, y = q.get()
            if x + 1 < len(matrix) and not exist[x+1][y]:
                q.put((matrix[x+1][y], x+1, y))
                exist[x+1][y] = True
            if y + 1 < len(matrix[0]) and not exist[x][y+1]:
                q.put((matrix[x][y+1], x, y+1))
                exist[x][y+1] = True
            k -= 1
        return cur