leetcode 373. Find K Pairs with Smallest Sums

原题链接

原题链接

解题思路

解这道题花了一点时间。发现用PriorityQueue来解决这问题就变得简单了。

1. 先用一个类封装了row，col，val，并实现compareTo方法，方便用PriorityQueue排序。
2. 先将最小值0,0,nums1[0]+nums2[0]压入优先队列。
3. 算法思想是bfs，从优先队列取出最小值并加入List res，很明显最接近这个值的要么（row+1，col），要么（row，col+1）。将两个值压入优先队列。
4. 返回res。

解题代码

public class Solution {

    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int x = nums1.length;int y = nums2.length;
        Queue<Pair> q = new PriorityQueue<>();
        List<int[]> res = new ArrayList<>();
        if (nums1 == null || nums2 == null || nums1.length == 0 || nums2.length == 0 || k == 0) {
            return res;
        }
        boolean[][] visited = new boolean[x][y];
        q.offer(new Pair(0,0,nums1[0]+nums2[0]));
        visited[0][0] = true;
        for(int i = 0;i<k;i++){
            if(q.isEmpty()) break;
            Pair p = q.poll();
            res.add(new int[]{nums1[p.row],nums2[p.col]});
            if(p.row < nums1.length - 1 && !visited[p.row+1][p.col]){
                q.offer(new Pair(p.row + 1,p.col,nums1[p.row+1]+nums2[p.col]));
                visited[p.row+1][p.col] = true;
            }
            if(p.col < nums2.length-1 && !visited[p.row][p.col+1]){
                q.offer(new Pair(p.row,p.col+1,nums1[p.row]+nums2[p.col+1]));
                visited[p.row][p.col+1] = true;
            }
        }
      return res;
    }
}
class Pair implements Comparable<Pair>{
    int row;
    int col;
    int val;

    Pair(int r,int c,int v){
        this.row = r;
        this.col = c;
        this.val = v;
    }
    @Override
    public int compareTo(Pair pair){
        return this.val - pair.val;
    }
}