leetcode 373. Find K Pairs with Smallest Sums

原题链接

原题链接

解题思路

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

先用一个类封装了row，col，val，并实现compareTo方法，方便用PriorityQueue排序。

先将最小值0,0,nums1[0]+nums2[0]压入优先队列。

算法思想是bfs，从优先队列取出最小值并加入List res，很明显最接近这个值的要么（row+1，col），要么（row，col+1）。将两个值压入优先队列。

返回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;
}
}