K Pairs with smallest sum in two arrays in Java

2 comments on “K Pairs with smallest sum in two arrays in Java”

• 
  

  Harshith M P 1SG20IS036

  import java.util.*;

  public class Main {

  public static void kSmallestPair(int[] A, int[] B,
  int K)
  {
  int n = A.length;
  // Min heap which contains tuple of the format
  // (sum, (i, j)) i and j are the indices
  // of the elements from array A
  // and array B which make up the sum.
  PriorityQueue pq
  = new PriorityQueue((a, b) -> a[0] – b[0]);

  // my_set is used to store the indices of
  // the pair(i, j) we use my_set to make sure
  // the indices does not repeat inside min heap.
  Set mySet = new HashSet();

  // initialize the heap with the minimum sum
  // combination i.e. (A[0] + B[0])
  // and also push indices (0,0) along
  // with sum.
  pq.offer(new int[] { A[0] + B[0], 0, 0 });
  mySet.add(“0,0”);

  // iterate upto K
  int count = 0;
  while (count < K && !pq.isEmpty()) {
  // array format (sum, i, j).
  int[] temp = pq.poll();
  int i = temp[1], j = temp[2];
  System.out.println(
  "(" + A[i] + ", " + B[j]
  + ")"); // Extracting pair with least sum
  // such that one element is from
  // arr1 and another is from arr2

  // check if i+1 is in the range of our first
  // array A
  boolean flag = false;
  if (i + 1 < n) {
  int sum = A[i + 1] + B[j];
  // insert (A[i + 1] + B[j], i + 1, j)
  // into min heap.
  String key = (i + 1) + "," + j;

  // insert only if the pair (i + 1, j) is
  // not already present inside the set i.e.
  // no repeating pair should be present
  // inside the heap.
  if (!mySet.contains(key)) {
  pq.offer(new int[] { sum, i + 1, j });
  mySet.add(key);
  flag = true;
  }
  }

  // check if j+1 is in the range of our second
  // array B
  if (j + 1 < B.length) {
  // insert (A[i] + B[j + 1], i, j + 1)
  // into min heap.
  int sum = A[i] + B[j + 1];
  String key = i + "," + (j + 1);

  // insert only if the pair (i, j + 1)
  // is not present inside the heap.
  if (!mySet.contains(key)) {
  pq.offer(new int[] { sum, i, j + 1 });
  mySet.add(key);
  flag = true;
  }
  }
  if (!flag) {
  break;
  }
  count++;
  }
  }

  public static void main(String[] args)
  {
  int[] A = { 1 };
  int[] B = { 2, 4, 5, 9 };
  int K = 8;
  kSmallestPair(A, B, K);
  }
  }

  Log in to Reply
• 
  

  PrepInsta Support

  Join our Discord for technical queries.

  Log in to Reply