Thinking process

The sum of the first element in each array should be the smallest.

similar to LC378

Mistakes

seems like adding the first pair [i, j] and then add [i, j + 1] and [i + 1, j] doesn't work

because in this way, we ween extra space to remember whether the pair has been visited before

• so, add the pair of each element in the first array and the first element in the second array
• then, move the index of the second array

Time Complexity

O(k logk)

class Pair {
    public int num1;
    public int num2;
    public int j;
    public int sum;

    public Pair (int num1, int num2, int j) {
        this.num1 = num1;
        this.num2 = num2;
        this.j = j;
        this.sum = num1 + num2;
    }
}

class PairComparator implements Comparator<Pair> {
    public int compare(Pair p1, Pair p2) {
        return p1.sum - p2.sum;
    }
}

public class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int n = nums1.length;
        int m = nums2.length;
        List<int[]> res = new ArrayList<>();
        if (n == 0 || m == 0 || k == 0) return res;

        PriorityQueue<Pair> pq = new PriorityQueue<Pair>(k, new PairComparator());

        for (int i = 0; i < n && i <= k; i++) {
            pq.add(new Pair(nums1[i], nums2[0], 0));
        }

        int count = 0;

        while (!pq.isEmpty() && count < k) {
            Pair p = pq.poll();

            if (p.j + 1 < m) {
                pq.add(new Pair(p.num1, nums2[p.j + 1], p.j + 1));
            }

            int[] ans = new int[]{p.num1, p.num2};
            res.add(ans);

            count++;
        } 


        return res;
    }
}