K Smallest Sum Pairs

$\mathtt{REFERENCE}$ @ LeetCode

$\mathtt{RELATED\ PROBLEM\ -\ Kth\ Smallest\ Numer\ in\ Sorted\ Matrix}$

Task Description

Given two integer arrays nums1 and nums1 sorted in non-decreasing order and an integer k.

Define a pair (u, v) consisting one element from the first array and one element from the second array.

Find k pairs with smallest sums.

Solution 1 : Brute Force Solution

The naive solution would be to find all possible pairs and find the best k pairs from it. If the two array contain n and m items, then there are nm possible pairs. Retrieving $k$ best pairs will take ${\cal \theta}(klog(mn))$ time.

Solution 2: Efficient Solution

This Brute Force solution can be optimized by using Priority Queue. The Priority Queue will hold sequence of number pairs (i, j) where $i$ is index in the first array and $j$ is that in the second array. Priority of the queue is determined by the sum of numbers at these indices.

  PriorityQueue<int[]> queue = new PriorityQueue<int[]>(new Comparator<int[]>(){
      @Override
      public int compare(int[] a, int[] b) {
          return nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]];
      }
  });

    

Initially, for $i\ =\ 0$, all pairs with the second array is inserted into the priority queue. At this time, the first priority element is minimum element.

  for (int i = 0; i < nums2.length; i++) {
      queue.offer(new int[]{0, i});
  }

    

Then we keep on polling an element from the queue and insert into the result. For each poll, we will insert into the queue, $i + 1$ pairing with its corresponding unpaired index so far in the second array. This is repeated untill the queue is not empty or untill we have all the $k$ pairs required.

  while(!queue.isEmpty() && res.size() < k) {
      int[] item = queue.poll();
      res.add(new int[]{nums1[item[0]], nums2[item[1]]});
      
      if (item[0] == nums1.length - 1) {
          continue;
      } else {
          queue.offer(new int[]{item[0] + 1, item[1]});
      }
  }

    

Full Code

    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> res = new ArrayList<int[]>();
        PriorityQueue<int[]> queue = new PriorityQueue<int[]>(new Comparator<int[]>(){
            @Override
            public int compare(int[] a, int[] b) {
                return nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]];
            }
        });
        
        if(nums1.length == 0 || nums2.length == 0) {
            return res;
        }
        
        for (int i = 0; i < nums2.length; i++) {
            queue.offer(new int[]{0, i});
        }
        
        while(!queue.isEmpty() && res.size() < k) {
            int[] item = queue.poll();
            res.add(new int[]{nums1[item[0]], nums2[item[1]]});
            
            if (item[0] == nums1.length - 1) {
                continue;
            } else {
                queue.offer(new int[]{item[0] + 1, item[1]});
            }
        }
    
        return res;
    }

    

Time Complexity

Initial insertion into the priority queue takes ${\cal O}(m)$ time. During the iteration, for each minimum element polled from the queue, we insert at maximum only one element, hence the size of the queue is maintained ${\cal O}(m)$ and each iteration take ${\cal O}(log(m))$ time. The iteration is repeated untill we have all the $k$ results or the queue is empty. $\therefore$ Overall Time complexity is ${\cal O}(k log(m)) \approx {\cal O}(k log(n))$. The space complexity is ${\cal O}(n)$ for the size of priority queue.

---

---

---