2.1 Find the 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.

Assumption: the matrix is not null, n, m > 0. 0 < k <= n*m

Input: int[][] matrix, int k

Output: int the smallest number

Solution:

1. Initial state: [0][0]

2. Expand[i][j] → generate [i][j+1]

generate [i+1][j]

3. Terminate condition: after expanding k times

4. Dedup: use boolean[][] to mark generated numbers

TC: k iterations to run k

for each iteration:

pop one element out of pq log2k (pq.size <= 2k)

generate right log2k

bottom log2k

Total: k * 3log2k = O(klogk)

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

public int kthSmallest(int[][] matrix, int k) {
  int rows = matrix.length;
  int cols = matri[0].length;
  PriorityQueue<Cell> minHeap = new PriorityQueue<Cell>(k, new Comparator<Cell>(){
    @Override
    public int compare(Cell c1, Cell c2) {
      if (c1.value == c2.value) { return 0; }
      return c1.value < c2.value ? -1 : 1;
    }
});

  boolean[][] visited = new boolean[rows][cols];
  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.col]) {
      minHeap.offer(new Cell(cur.row + 1, cur.col, matrix[cur.row + 1][cur.col]));
      visited[cur.row + 1][cur.col] = true;
    }
    if (cur.col + 1< cols && !visited[cur.row][cur.col + 1]) {
       minHeap.offer(new Cell(cur.row, cur.col + 1, matrix[cur.row][cur.col + 1]));
       visited[cur.row][cur.col + 1] = true;
    }
  }
  return minHeap.peek().value;
}

Similar solution to question: Kth Smallest Sum In Two Sorted Arrays

Given two sorted arrays A and B, of sizes m and n respectively. Define sum = a + b, where a is one element from A and b is one element from B. Find the Kth smallest sum out of all possible sums. Assumption: A and B are not null and length > 0. k > 0 and K <= m * n.

class Cell {
   int i;
   int j;
   int sum;
   public Cell(int i, int j, int sum) {
     this.i = i;
     this.j = j;
     this.sum = sum;
   }
 }

 public int kthSum(int[] A, int[] B, int k) {
   PriorityQueue<Cell> pq = new PriorityQueue<>(1, new Comparator<Cell>(){
     @Override
     public int compare(Cell a, Cell b) {
       return a.sum - b.sum;
     }
   });
   boolean[][] visited = new boolean[A.length][B.length];
   pq.offer(new Cell(0, 0, A[0] + B[0]));
   visited[0][0] = true;
   for (int i = 1; i < k; i++) {
     Cell cur = pq.poll();
     if (cur.i + 1 < A.length && !visited[cur.i + 1][cur.j]) {
       pq.offer(new Cell(cur.i + 1, cur.j, A[cur.i + 1] + B[cur.j]));
       visited[cur.i + 1][cur.j] = true;
     }
     if (cur.j + 1 < B.length && !visited[cur.i][cur.j + 1]) {
       pq.offer(new Cell(cur.i, cur.j + 1, A[cur.i] + B[cur.j + 1]));
       visited[cur.i][cur.j + 1] = true;
     }
   }
   return pq.peek().sum;
 }