原题
解
第一步，万年不变的查错。如果给的array是null或空，或者k等于0，那么直接return。
public int kthSmallestSum(int[] A, int[] B, int k) {
if (A == null || A.length == 0 || B == null || B.length == 0 || k == 0) {
return 0;
}
...
}
思路大概跟排序矩阵中的从小到大第k个数一样，做一个MinHeap，存放最多2k个数，每次poll最小的数，然后下一个最小的只可能在(i+1, j)和(i, j+1)中出现。把这两个再放进heap中，当然要去重，然后再取最小，继续k-1次。剩下的当前root就是第k小的。
第一步先建一个Pair class用来存放坐标。这里需要override equals 和hashCode，因为一会去重要把这些扔进HashSet，所以需要hashCode，然后如果hashCode重复，要通过equals来判断两个Pair是不是一样的。这里yong
private class Pair {
int a;
int b;
int sum;
public Pair(int a, int b, int sum) {
this.a = a;
this.b = b;
this.sum = val;
}
@Override
public boolean equals(Object other) {
if (!(other instanceof Pair)) {
return false;
}
Pair p = (Pair) other;
return this.a == p.a && this.b == p.b;
}
@Override
public int hashCode() {
return a * 31 + b;
}
}
然后初始化需要的东西。两个arrayda,db用来存放下一个位置的变量。一个MinHeap用来找最小的数。一个HashSet用来去重。MinHeap加入第一个最小的和，即A里面最小的数加B里面最小的数。第一个不需要放到HashSet里面， 因为我们找下一个数都是i+1或者j+1，所以从(0, 0)开始，是不可能有下一个回到(0, 0)的。
int[] da = new int[] {0, 1};
int[] db = new int[] {1, 0};
Queue minHeap = new PriorityQueue<>(Comparator.comparingInt(a -> a.sum));
Set hash = new HashSet<>();
minHeap.offer(new Pair(0, 0, A[0] + B[0]));
然后for循环k-1次，即从1到k-1。每次取出最小数，然后找下面两个可能的最小和，去重，放进heap。
for (int i = 1; i < k; i++) {
Pair current = minHeap.poll();
for (int j = 0; j < 2; j++) {
int nextA = current.a + da[j];
int nextB = current.b + db[j];
if (nextA < 0 || nextA >= A.length || nextB < 0 || nextB >= B.length) {
continue;
}
Pair pair = new Pair(nextA, nextB, A[nextA] + B[nextB]);
if (!hash.contains(pair)) {
minHeap.offer(pair);
hash.add(pair);
}
}
}
最后返回MinHeap的第一个即可，因为循环找出了k-1个最小的和，所以剩下最小的就是第k小的。
return minHeap.peek().sum;
完整的code
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 Pair {
int a;
int b;
int sum;
public Pair(int a, int b, int sum) {
this.a = a;
this.b = b;
this.sum = sum;
}
@Override
public boolean equals(Object other) {
if (!(other instanceof Pair)) {
return false;
}
Pair p = (Pair) other;
return this.a == p.a && this.b == p.b;
}
@Override
public int hashCode() {
return a * 31 + b;
}
}
public int kthSmallestSum(int[] A, int[] B, int k) {
if (A == null || A.length == 0 || B == null || B.length == 0 || k == 0) {
return 0;
}
int[] da = new int[] {0, 1};
int[] db = new int[] {1, 0};
Queue minHeap = new PriorityQueue<>(Comparator.comparingInt(a -> a.sum));
Set hash = new HashSet<>();
minHeap.offer(new Pair(0, 0, A[0] + B[0]));
for (int i = 1; i < k; i++) {
Pair current = minHeap.poll();
for (int j = 0; j < 2; j++) {
int nextA = current.a + da[j];
int nextB = current.b + db[j];
if (nextA < 0 || nextA >= A.length || nextB < 0 || nextB >= B.length) {
continue;
}
Pair pair = new Pair(nextA, nextB, A[nextA] + B[nextB]);
if (!hash.contains(pair)) {
minHeap.offer(pair);
hash.add(pair);
}
}
}
return minHeap.peek().sum;
}
};
分析
空间复杂度是O(2k)，即O(k)。因为要放最多2k个数字进入heap。然后poll了k次，每次算log(k)，那时间复杂度就是O(klogk)。
同排序矩阵中的从小到大第k个数一样，这个题还可以用二分法，也是适用于数很多，但range很小的时候。
总的来说，这道题就是排序矩阵第k小的题的马甲。解法完全一样，i和j坐标被分成了两个array，然后数值成为了两个array的和。