Skip to content

Instantly share code, notes, and snippets.

@kenkoooo kenkoooo/Solution.java Secret
Created Dec 31, 2019

Embed
What would you like to do?
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;
import java.util.TreeSet;
class Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
List<List<Integer>> result = new ArrayList<>();
if (nums1.length == 0 || nums2.length == 0 || k == 0) {
return result;
}
PriorityQueue<int[]> Q = new PriorityQueue<>((a, b) -> nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]]);
TreeSet<int[]> added = new TreeSet<>((a, b) -> {
if (a[0] != b[0]) {
return a[0] - b[0];
}
return a[1] - b[1];
});
Q.add(new int[]{0, 0});
added.add(new int[]{0, 0});
while (result.size() < k && !Q.isEmpty()) {
int[] ans = Q.poll();
int i = ans[0];
int j = ans[1];
result.add(Arrays.asList(nums1[i], nums2[j]));
if (i + 1 < nums1.length && !added.contains(new int[]{i + 1, j})) {
Q.add(new int[]{i + 1, j});
added.add(new int[]{i + 1, j});
}
if (j + 1 < nums2.length && !added.contains(new int[]{i, j + 1})) {
Q.add(new int[]{i, j + 1});
added.add(new int[]{i, j + 1});
}
}
return result;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.