Java heap Solution 13ms

• Use a heap, first put i = 0, j = 0 into heap. Once a pair (current min) is polled from heap, add (i + 1, j) and (i, j + 1) into heap if they have not been added.

``````public class Solution {

private class Node implements Comparable<Node>{
int i, j;
int val;
public Node(int i, int j, int val) {
this.i = i;
this.j = j;
this.val = val;
}
@Override
public int compareTo(Node that) {
return this.val - that.val;
}
}

public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
Queue<Node> q = new PriorityQueue<>();
List<int[]> res = new ArrayList<>();
int m = nums1.length, n = nums2.length;
if (m == 0 || n == 0) return res;
Set<Integer> set = new HashSet<>();
q.offer(new Node(0, 0, nums1[0] + nums2[0]));
for (int t = 0; t < k; t++) {
if (q.isEmpty()) break;
Node x = q.poll();
if (x.i < nums1.length - 1 && !set.contains(x.i * n + n + x.j)) {
q.offer(new Node(x.i + 1, x.j, nums1[x.i + 1] + nums2[x.j]));
set.add(x.i * n + n + x.j);
}
if (x.j < nums2.length - 1 && !set.contains(x.i * n + x.j + 1)) {
q.offer(new Node(x.i, x.j + 1, nums1[x.i] + nums2[x.j + 1]));
set.add(x.i * n + x.j + 1);
}
}
return res;
}
}
``````

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.