Find K Pairs with Smallest Sum

 1 public class Solution {
 2     private Set<Integer> visited;
 3     public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
 4         List<int[]> result = new ArrayList<>();
 5         visited = new HashSet<>();
 6         PriorityQueue<int[]> queue = new PriorityQueue<>(new Comparator<int[]>() {
 7             @Override
 8             public int compare(int[] a, int[] b) {
 9                 return (nums1[a[0]] + nums2[a[1]]) - (nums1[b[0]] + nums2[b[1]]);
10             }
11         });
12         add(queue, nums1, nums2, 0, 0);
13         while (queue.size() > 0 && result.size() < k) {
14             int[] current = queue.poll();
15             result.add(new int[]{nums1[current[0]], nums2[current[1]]});
16             add(queue, nums1, nums2, current[0] + 1, current[1]);
17             add(queue, nums1, nums2, current[0], current[1] + 1);
18         }
19         return result;
20     }
21     
22     private void add(PriorityQueue<int[]> queue, int[] nums1, int[] nums2, int n1, int n2) {
23         if (n1 < nums1.length && n2 < nums2.length && !visited.contains(n1 * nums2.length + n2)) {
24             visited.add(n1 * nums2.length + n2);
25             queue.offer(new int[] {n1, n2});
26         }
27     }
28 }
Dijkstra Algorithm
posted on 2016-07-08 05:55 keepshuatishuati 阅读(166) 评论(0) 编辑收藏举报