373. Find K Pairs with Smallest Sums && 378. Kth Smallest Element in a Sorted Matrix

373. Find K Pairs with Smallest Sums

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u,v) which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1),(u2,v2) ...(uk,vk) with the smallest sums.

Example 1:

Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3
Return: [1,2],[1,4],[1,6]
The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2
Return: [1,1],[1,1]
The first 2 pairs are returned from the sequence:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Given nums1 = [1,2], nums2 = [3],  k = 3 
Return: [1,3],[2,3]
All possible pairs are returned from the sequence:
[1,3],[2,3]

 Heap

 

This problem is almost the same as problem 378.

Here I gave two solutions as for 378.

Solution 1: Doing repeatitive work without using heap.

public class Solution {
  public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
    List<int[]> ret = new ArrayList<>();
    if (nums1.length == 0 || nums2.length == 0 || k == 0)
      return ret;
    //Each number in nums1 will pair with one number in nums2.
    //The index of pair number in nums2 in kept in the array num2Index.
    int[] num2Index = new int[nums1.length];
    int num1start = 0;
    while (k-- > 0) {
      //every time we find a min, we will restart from the first number in nums1.
      int currentMin = Integer.MAX_VALUE;
      int nums1MinIndex = -1;
      for (int iIn1 = num1start; iIn1 < nums1.length; ++iIn1) {
        if (num2Index[iIn1] >= nums2.length) {//nums[iIn1] has been fully paired.
          num1start = iIn1; // we only need to start to iIn1 for the next min search
          continue;
        }
        int nextMin = nums1[iIn1] + nums2[num2Index[iIn1]];
        if (nextMin < currentMin) {
          currentMin = nextMin;
          nums1MinIndex = iIn1;
        }
      }
      if (nums1MinIndex == -1)
        return ret; //cannot find any more pairs, which mean k is larger than total combination counts
      ret.add(new int[]{nums1[nums1MinIndex], nums2[num2Index[nums1MinIndex]]});
      num2Index[nums1MinIndex]++; //set the corresponding index to next one
    }
    return ret;
  }
}

 

Solution 2: Algorithm based on heap, better performance.

public class Solution {
  public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
    List<int[]> ret = new ArrayList<>();
    int L1 = nums1.length;
    int L2 = nums2.length;
    if (L1 == 0 || L2 == 0 || k == 0)
      return ret;

    PriorityQueue<Tuple> queue = new PriorityQueue<>();
    for (int i = 0; i < L1; ++i)
      queue.add(new Tuple(i, 0, nums1[i] + nums2[0]));

    while (k-- > 0 && !queue.isEmpty()) {
      Tuple nextSmall = queue.poll();
      ret.add(new int[]{nums1[nextSmall.index1], nums2[nextSmall.index2]});
      if (nextSmall.index2 < L2 - 1)
        queue.offer(new Tuple(nextSmall.index1, nextSmall.index2 + 1,
          nums1[nextSmall.index1] + nums2[nextSmall.index2 + 1]));
    }
    return ret;
  }
}

class Tuple implements Comparable<Tuple> {
  int index1, index2, sum;

  Tuple(int index1, int index2, int sum) {
    this.index1 = index1;
    this.index2 = index2;
    this.sum = sum;
  }

  @Override
  public int compareTo(Tuple that) {
    return this.sum - that.sum;
  }
}

 

378. Kth Smallest Element in a Sorted Matrix

Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Example:

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,
return 13

Note: You may assume k is always valid, 1 ≤ k ≤ n2.

 Binary Search Heap

 

First thought:

Compare the smallest number for each row, and return the smallest one. But we are doing repeatitive work for each round. Use priority queue to save the work. See next section.

public class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int ROW = matrix.length;
        int COL = matrix[0].length;
        //use an array to keep the column index for each row.
        int[] columnIndices = new int[ROW];
        
        int nextSmallRow = 0;
        int nextSmall = Integer.MAX_VALUE;
        int startRow = 0;
        while(k>=1) {
            for(int r=startRow; r<ROW; ++r) {
                int c = columnIndices[r];
                if(c == COL){
                    startRow = r+1;
                    continue;
                }
                int current = matrix[r][c];
                if(current<nextSmall) {
                    nextSmall = current;
                    nextSmallRow = r;
                }
                if(columnIndices[r] == 0) //this is the last comparison for this round.
                    break;
            }
            if(k==1)
                return nextSmall;
            --k;
            ++columnIndices[nextSmallRow];
            nextSmall = Integer.MAX_VALUE;//reset nextSmall, and start next round
        }
        return -999;
    }
}

 

Same idea, but based on priority queue to save comparison work.

class Solution {
  public int kthSmallest(int[][] matrix, int k) {
    int ROW = matrix.length;
    int COL = matrix[0].length;
    PriorityQueue<Tuple> queue = new PriorityQueue<>();
    for (int r = 0; r < ROW; ++r)
      queue.offer(new Tuple(r, 0, matrix[r][0])); //add first element of each row to queue.

    Tuple nextSmall = null;
    while (k > 0) {
      nextSmall = queue.poll();
      --k;//we've found one, nextSmall
      if (nextSmall.col < COL-1)//if not reach the end of this row, add next one to queue
        queue.offer(new Tuple(nextSmall.row, nextSmall.col + 1, matrix[nextSmall.row][nextSmall.col + 1]));
    }
    return nextSmall.val;
  }
}

class Tuple implements Comparable<Tuple> {
  int row, col, val;

  Tuple(int row, int col, int val) {
    this.row = row;
    this.col = col;
    this.val = val;
  }

  @Override
  public int compareTo(Tuple that) {
    return this.val - that.val;
  }
}