Program to find the minimum range of set that includes at least one element from each of the k sorted lists

MinimumRange.java

/**
* @author Sabbir Manandhar
*         ManandharSabbir@gmail.com
*/
class MinimumRange {
    public int[] smallestRange(List<List<Integer>> nums) {
        // Min heap to store intermediate k elements
        PriorityQueue<int[]> heap = new PriorityQueue<>(new Comparator<int[]>() {
            public int compare(int[] a, int[] b) {
                return a[0] - b[0];
            }
        });
        // keep track of max in the heap to compute range
        int maxSoFar = Integer.MIN_VALUE;
        
        // start with minimum elements from the lists into our heap
        for(int i = 0; i < nums.size(); i++) {
            int num = nums.get(i).get(0);
            heap.offer(new int[]{num, i, 0});
            maxSoFar = Math.max(maxSoFar, num);
        }
        
        // keep track of minimum range
        int rangeS = heap.peek()[0];
        int rangeE = maxSoFar;
        
        while(true) {
            // remove minimum element from the heap
            int[] min = heap.poll();
            int nextIdx = min[2] + 1;
          
            // once one of the list exhausts return
            if(nextIdx == nums.get(min[1]).size()) {
                break;
            }
          
            // insert next element fromt the list
            int nextNum = nums.get(min[1]).get(nextIdx);
            heap.offer(new int[]{nextNum, min[1], nextIdx});
            
            // update max so far and minimum range
            maxSoFar = Math.max(maxSoFar, nextNum);
            int newRangeS = heap.peek()[0];
            int newRangeE = maxSoFar;
            
            if (newRangeE - newRangeS < rangeE - rangeS) {
                rangeS = newRangeS;
                rangeE = newRangeE;
            }
        }
        return new int[]{rangeS, rangeE};
    }
}