classSolution { booleanflag=true; public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) { List<List<Integer>> ans = newArrayList<>(); intn= nums1.length, m = nums2.length; if (n > m && !(flag = false)) return kSmallestPairs(nums2, nums1, k); PriorityQueue<int[]> q = newPriorityQueue<>((a,b)->(nums1[a[0]]+nums2[a[1]])-(nums1[b[0]]+nums2[b[1]])); for (inti=0; i < Math.min(n, k); i++) q.add(newint[]{i, 0}); while (ans.size() < k && !q.isEmpty()) { int[] poll = q.poll(); inta= poll[0], b = poll[1]; ans.add(newArrayList<>(){{ add(flag ? nums1[a] : nums2[b]); add(flag ? nums2[b] : nums1[a]); }}); if (b + 1 < m) q.add(newint[]{a, b + 1}); } return ans; } }
Python3 代码:
1 2 3 4 5 6 7 8 9 10 11 12 13 14
class Solution: def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]: flag, ans = (n := len(nums1)) > (m := len(nums2)), [] if flag: n, m, nums1, nums2 = m, n, nums2, nums1 pq = [] for i in range(min(n, k)): heapq.heappush(pq, (nums1[i] + nums2[0], i, 0)) while len(ans) < k and pq: _, a, b = heapq.heappop(pq) ans.append([nums2[b], nums1[a]] if flag else [nums1[a], nums2[b]]) if b + 1 < m: heapq.heappush(pq, (nums1[a] + nums2[b + 1], a, b + 1)) return ans
funckSmallestPairs(nums1 []int, nums2 []int, k int) [][]int { n, m, ans := len(nums1), len(nums2), [][]int{} flag := n > m if flag { n, m, nums1, nums2 = m, n, nums2, nums1 } if n > k { n = k } pq := make(hp, n) for i := 0; i < n; i++ { pq[i] = []int{nums1[i] + nums2[0], i, 0} } heap.Init(&pq) for pq.Len() > 0 && len(ans) < k { poll := heap.Pop(&pq).([]int) a, b := poll[1], poll[2] if flag{ ans = append(ans, []int{nums2[b], nums1[a]}) }else{ ans = append(ans, []int{nums1[a], nums2[b]}) } if b < m - 1 { heap.Push(&pq, []int{nums1[a] + nums2[b + 1], a, b + 1}) } } return ans } // 最小堆模板 type hp [][]int func(h hp) Len() int { returnlen(h) } func(h hp) Less(i, j int) bool { return h[i][0] < h[j][0] } func(h hp) Swap(i, j int) { h[i], h[j] = h[j], h[i] } func(h *hp) Push(v interface{}) { *h = append(*h, v.([]int)) } func(h *hp) Pop() interface{} { a := *h; v := a[len(a)-1]; *h = a[:len(a)-1]; return v }
classSolution { int[] nums1, nums2; int n, m; public List<List<Integer>> kSmallestPairs(int[] n1, int[] n2, int k) { nums1 = n1; nums2 = n2; n = nums1.length; m = nums2.length; List<List<Integer>> ans = newArrayList<>(); intl= nums1[0] + nums2[0], r = nums1[n - 1] + nums2[m - 1]; while (l < r) { intmid= (int)(0L + l + r >> 1); if (check(mid, k)) r = mid; else l = mid + 1; } intx= r; for (inti=0; i < n; i++) { for (intj=0; j < m; j++) { if (nums1[i] + nums2[j] < x) { List<Integer> temp = newArrayList<>(); temp.add(nums1[i]); temp.add(nums2[j]); ans.add(temp); } elsebreak; } } for (inti=0; i < n && ans.size() < k; i++) { inta= nums1[i], b = x - a; intc= -1, d = -1; l = 0; r = m - 1; while (l < r) { intmid= (int)(0L + l + r >> 1); if (nums2[mid] >= b) r = mid; else l = mid + 1; } if (nums2[r] != b) continue; c = r; l = 0; r = m - 1; while (l < r) { intmid= (int)(0L + l + r + 1) >> 1; if (nums2[mid] <= b) l = mid; else r = mid - 1; } d = r; for (intp= c; p <= d && ans.size() < k; p++) { List<Integer> temp = newArrayList<>(); temp.add(a); temp.add(b); ans.add(temp); } } return ans; } booleancheck(int x, int k) { intans=0; for (inti=0; i < n && ans < k; i++) { for (intj=0; j < m && ans < k; j++) { if (nums1[i] + nums2[j] <= x) ans++; elsebreak; } } return ans >= k; } }