Kth Smallest Sum In Two Sorted Arrays 465
Question
Given two integer arrays sorted in ascending order and an integer k. Define sum = a + b, where a is an element from the first array and b is an element from the second one. Find the kth smallest sum out of all possible sums.
Example
Given [1, 7, 11] and [2, 4, 6].
For k = 3, return 7.
For k = 4, return 9.
For k = 8, return 15.
Challenge
Do it in either of the following time complexity:
O(k log min(n, m, k)). where n is the size of A, and m is the size of B.
O( (m + n) log maxValue). where maxValue is the max number in A and B.
Solution
idea:位置为i(A中),j(B中)元素(表示A[i]+B[j])的下一个元素为位置i+1,j和i,j+1之中的一个。
将A[0]和B[0]的和加入一个最小堆,并将位置0,0加入一个hashset记录加入过堆的位置
将最堆顶元素出队,得到其位置i,j,并将A[i+1]+B[j]和A[i]+B[j+1]加入堆中,其中i+1和j+1必须是合法的,即不越界,同时i+1,j和i,j+1没有加入过堆
重复2,出队k次即可
代码如下:
public class Solution {
/**
* @param A an integer arrays sorted in ascending order
* @param B an integer arrays sorted in ascending order
* @param k an integer
* @return an integer
*/
class Node implements Comparable<Node>{
int i;
int 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 another){
if(this.val == another.val){
return 0;
}
return this.val > another.val? 1 : -1;
}
}
int[] di = new int[]{1, 0};
int[] dj = new int[]{0, 1};
public int kthSmallestSum(int[] A, int[] B, int k) {
// Write your code here
if(A == null || B == null || A.length == 0 && B.length == 0 || k > A.length * B.length){
return 0;
}else if(A.length == 0){
return B[k];
}else if(B.length == 0){
return A[k];
}
HashSet<String> set = new HashSet<String>();
PriorityQueue<Node> queue = new PriorityQueue<Node>();
Node temp = new Node(0, 0, A[0] + B[0]);
queue.offer(temp);
set.add(temp.i + "," + temp.j);
Node next;
int i = 0;
int j = 0;
for(int n = 0; n < k - 1; n++){
temp = queue.poll();
for(int m = 0; m < 2; m++){
i = temp.i + di[m];
j = temp.j + dj[m];
next = new Node(i, j, 0);
if(i >= 0 && i < A.length && j >= 0 && j < B.length && !set.contains(next.i + "," + next.j)){
next.val = A[i] + B[j];
queue.offer(next);
set.add(next.i + "," + next.j);
}
}
}
return queue.peek().val;
}
}