Kth Smallest Number in sorted matrix

Given a matrix of size N x M. For each row the elements are sorted in ascending order, and for each column the elements are also sorted in ascending order. Find the Kth smallest number in it.

Assumptions

• the matrix is not null, N > 0 and M >0
• K >0 and K <= N * M

Examples

{ {1, 3, 5, 7},

{2, 4, 8, 9},

{3, 5, 11, 15},

{6, 8, 13, 18} }

• the 5th smallest number is 4
• the 8th smallest number is 6

public class Solution {
  class Cell {
    int row;
    int column;
    int val;

    public Cell(int row, int column, int val) {
      this.row = row;
      this.column = column;
      this.val = val;
    }
  }

  public int kthSmallest(int[][] matrix, int k) {
    // Write your solution here.
    int rows = matrix.length;
    int columns = matrix[0].length;

    PriorityQueue<Cell> minHeap = new PriorityQueue<Cell>(k, new Comparator<Cell>() {
      @Override
      public int compare(Cell c1, Cell c2) {
        if (c1.val == c2.val) {
          return 0;
          }
        return c1.val < c2.val ? -1 : 1;
        }
      });

    boolean[][] visited = new boolean[rows][columns];
    minHeap.offer(new Cell(0, 0, matrix[0][0]));
    visited[0][0] = true;


    for (int i = 0; i < k - 1; i++) {
      Cell cur = minHeap.poll();
      if (cur.row + 1 < rows && !visited[cur.row + 1][cur.column]) {
        minHeap.offer(new Cell(cur.row + 1, cur.column, matrix[cur.row + 1][cur.column]));
        visited[cur.row + 1][cur.column] = true;
      }
      if (cur.column + 1 < columns && !visited[cur.row][cur.column + 1]) {
        minHeap.offer(new Cell(cur.row, cur.column + 1, matrix[cur.row][cur.column + 1]));
        visited[cur.row][cur.column + 1] = true;
      }
    }
    return minHeap.peek().val;
  }
}
// 首先语法code能力还不太熟。找第k个最小值，想到bfs2（起点到终点的最小花费）
//终点就是第k个点。还用到了minHeap的性质， pop出来的是单调递增的。 pop k－1 次后堆顶就是答案
//用priorityQueue 数据结构来存。 expand了终点就结束了。
//Time: klog(k). pop k次，offer 2k次 3klogk