[Algo] 26. 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 {
  public int kthSmallest(int[][] matrix, int k) {
    // Write your solution here
    int row = matrix.length;
    int col = matrix[0].length;
    boolean[][] isVisited = new boolean[row][col];
    PriorityQueue<Cell> pq = new PriorityQueue<>(k, new Comparator<Cell>() {
      @Override
      public int compare(Cell a, Cell b) {
        return a.value - b.value;
      }
    });
    pq.offer(new Cell(0, 0, matrix[0][0]));
    isVisited[0][0] = true;
    for (int i = 0; i < k - 1; i++) {
      Cell cur = pq.poll();
      int curRow = cur.row;
      int curCol = cur.col;
      if (curRow + 1 < row && !isVisited[curRow + 1][curCol]) {
        pq.offer(new Cell(curRow + 1, curCol, matrix[curRow + 1][curCol]));
        isVisited[curRow + 1][curCol] = true;
      }
      if (curCol + 1 < col && !isVisited[curRow][curCol + 1]) {
        pq.offer(new Cell(curRow, curCol + 1, matrix[curRow][curCol + 1]));
        isVisited[curRow][curCol + 1] = true;
      }
    }
    return pq.poll().value;
  }
}

class Cell {
  int row;
  int col;
  int value;
  public Cell(int row, int col, int value) {
    this.row = row;
    this.col = col;
    this.value = value;
  }
}