Kth Smallest Sum In Two Sorted Arrays

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:

1. O(k log min(n, m, k)). where n is the size of A, and m is the size of B.
2. O( (m + n) log maxValue). where maxValue is the max number in A and B.

Solution:

Similar to Kth Smallest Number in Sorted Matrix. PriorityQueue with a sze min(n, m, k) is used. We put (1,2) into PriorityQueue and then poll the smallest sum out (which is (1,2) itself). After that, we add the neighbors of (1,2) into PriorityQueue (1,4 and 7,2). Repeat these steps until finding k smallest sum.

Code:

class SumType {
    int a; //a idx
    int b; //b idx
    int sum; // sum of A[a] and B[b]
    public SumType(int inA, int inB, int inSum) {
        this.a = inA;
        this.b = inB;
        this.sum = inSum;
    }
}

public class Solution {

    public int kthSmallestSum(int[] A, int[] B, int k) {
        if (A == null || B == null) {
            return Integer.MIN_VALUE;
        }
        int n = A.length;
        int m = A.length;
        if (n * m < k) {
            return Integer.MIN_VALUE;
        }
        Queue<SumType> q = new PriorityQueue(Math.min(Math.min(n, m), k),
            new Comparator<SumType>() {
                public int compare(SumType a, SumType b) {
                    return a.sum - b.sum;
                }
        });
        HashSet<Integer> visited = new HashSet<>();
        kthSmallestSumHelper(0, 0, A, B, q, visited);
        int result = Integer.MIN_VALUE;
        for (int i = 0; i < k; i++) {
            SumType temp = q.poll();
            result = temp.sum;
            kthSmallestSumHelper(temp.a + 1, temp.b, A, B, q, visited);
            kthSmallestSumHelper(temp.a, temp.b + 1, A, B, q, visited);
        }
        return result;
    }

    private void kthSmallestSumHelper(int a, int b, int[] A, int[] B,
                    Queue<SumType> q, HashSet<Integer> visited) {
        if (!(a < A.length && b < B.length)) {
            return;
        }
        int hashIdx = a * B.length + b;
        if (!visited.contains(hashIdx)) {
            q.offer(new SumType(a, b, A[a] + B[b]));
            visited.add(hashIdx);
        }
    }
}