这道题特别有意思
本质上就是Kth Smallest Element in a Sorted Matrix
但是很难想到
做法都是一样的
写了一个bug就是在加入a+1和b+1的时候要提前check有没有越界,越界了就不能再加入了
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
*/
private class Tuple {
public int a, b;
public int sum;
public Tuple(int a, int b, int sum) {
this.a = a;
this.b = b;
this.sum = sum;
}
}
public int kthSmallestSum(int[] A, int[] B, int k) {
// Write your code here
if (A == null || A.length == 0 || B == null || B.length == 0 || k <= 0) throw new IllegalArgumentException();
//min heap
PriorityQueue<Tuple> pq = new PriorityQueue<>(k, new Comparator<Tuple>() {
public int compare(Tuple t1, Tuple t2) {
return t1.sum - t2.sum;
}
});
pq.add(new Tuple(0, 0, A[0] + B[0]));
int count = k;
Tuple curr = null;
while (count > 0) {
curr = pq.poll();
count--;
//System.out.println("a:" + curr.a + " b:" + curr.b + " count:" + count);
if (curr.a == 0) {
if(curr.b + 1 < B.length) pq.offer(new Tuple(curr.a, curr.b + 1, A[curr.a] + B[curr.b + 1]));
}
if (curr.a + 1 < A.length) pq.offer(new Tuple(curr.a + 1, curr.b, A[curr.a + 1] + B[curr.b]));
}
return curr.sum;
}
}
No comments:
Post a Comment