Java Code using PriorityQueue. similar to merge k array


  • 26
    X

    Image you are merging k sorted array using a heap. Then everytime you pop the smallest element out and add the next element of that array to the heap. By keep doing this, you will have the smallest range.

    public int[] smallestRange(int[][] nums) {
    		PriorityQueue<Element> pq = new PriorityQueue<Element>(new Comparator<Element>() {
    			public int compare(Element a, Element b) {
    				return a.val - b.val;
    			}
    		});
    		int min = Integer.MAX_VALUE, max = Integer.MIN_VALUE;
    		for (int i = 0; i < nums.length; i++) {
    			Element e = new Element(i, 0, nums[i][0]);
    			pq.offer(e);
    			max = Math.max(max, nums[i][0]);
    		}
    		int range = Integer.MAX_VALUE;
    		int start = -1, end = -1;
    		while (pq.size() == nums.length) {
    
    			Element curr = pq.poll();
    			if (max - curr.val < range) {
    				range = max - curr.val;
    				start = curr.val;
    				end = max;
    			}
    			if (curr.idx + 1 < nums[curr.row].length) {
    				curr.idx = curr.idx + 1;
    				curr.val = nums[curr.row][curr.idx];
    				pq.offer(curr);
    				if (curr.val > max) {
    					max = curr.val;
    				}
    			}
    		}
    
    		return new int[] { start, end };
    	}
    
    	class Element {
    		int val;
    		int idx;
    		int row;
    
    		public Element(int r, int i, int v) {
    			val = v;
    			idx = i;
    			row = r;
    		}
    	}
    

  • 6
    T

    Based on your idea, here is my a little simpler code using int[] instead of Class Element

    public int[] smallestRange(List<List<Integer>> nums) {
    	PriorityQueue<int[]> minHeap = new PriorityQueue<int[]>(nums.size(), new Comparator<int[]>(){
    		public int compare(int[] o1, int[] o2) {
    			return o1[0] - o2[0];
    		}
    	});
    
    	int max = nums.get(0).get(0);
    	for(int i=0; i<nums.size(); i++) {
    		minHeap.add(new int[]{nums.get(i).get(0), i, 0});
    		max = Math.max(max, nums.get(i).get(0));
    	}
    	
    	int minRange = Integer.MAX_VALUE;
    	int start = -1;
    	while(minHeap.size() == nums.size()) {
    		int[] t = minHeap.poll();
    		if(max - t[0] < minRange) {
    			minRange = max - t[0];
    			start = t[0];
    		}
    		
    		if(t[2]+1 < nums.get(t[1]).size()) {
    			t[0] = nums.get(t[1]).get(t[2]+1);
    			t[2] ++;
    			minHeap.add(t);
    			max = Math.max(max, t[0]);
    		}
    	}
    	
    	return new int[]{start, start+minRange};
    }

  • 3
    Z

    Thanks for the solution! Here is my C++ version.

    class Solution {
    public:
        struct mycompare {
            bool operator () (pair<int, int>& a, pair<int, int>& b) {return a.first > b.first;}  
        };
        vector<int> smallestRange(vector<vector<int>>& nums) {
            int n = nums.size(), large = INT_MIN, maxlen = INT_MAX;
            priority_queue<pair<int, int>, vector<pair<int, int>>, mycompare> pq;
            for (int i = 0; i < nums.size(); i++) {
                large = max(large, nums[i][0]);
                pq.push({nums[i][0], i});
            }
            vector<int> ans(2, 0), idx(n, 0);
            while (pq.size() == n) {
                int cur = pq.top().first, row = pq.top().second;
                pq.pop();
                if (large-cur < maxlen) {
                    maxlen = large-cur;
                    ans[0] = cur;
                    ans[1] = large;
                }
                if (++idx[row] < nums[row].size()) {
                    pq.push({nums[row][idx[row]], row});
                    large = max(large, nums[row][idx[row]]);
                }
            }
            return ans;
        }
    };
    

  • 1
    S

    Python version.

    import heapq
    class Solution(object):
        def smallestRange(self, nums):
            """
            :type nums: List[List[int]]
            :rtype: List[int]
            """
            q = []  # element in the heap: val, i, j, where val is nums[i][j]
            max_val = nums[0][0]
            for i in range(len(nums)):
                heapq.heappush(q, (nums[i][0], i, 0))
                max_val = max(max_val, nums[i][0])  # also remember max of the heap
            min_range = [-10 ** 5, 10 ** 5]
            while q:
                min_val, i, j = heapq.heappop(q)
                if max_val - min_val < min_range[1] - min_range[0] or (
                                    max_val - min_val == min_range[1] - min_range[0] and min_val < min_range[0]):
                    min_range = [min_val, max_val]
                # push the next value in the ith array if any
                if j + 1 < len(nums[i]):
                    max_val = max(max_val, nums[i][j + 1])
                    heapq.heappush(q, (nums[i][j + 1], i, j + 1))
                else:  # ths ith array is exhausted
                    return min_range

  • 0
    N

    c# version

    public class Solution {
        public int[] SmallestRange(IList<IList<int>> nums)
        {
            var start = -1;
            var end = -1;
            var max = int.MinValue;
            var pq = new PriorityQueue<int[]>((a, b) => a[0] - b[0]);
            var range = int.MaxValue;
            for (var i = 0; i < nums.Count; i++)
            {
                pq.Enqueue(new[] { nums[i][0], i, 0 }); // val, row, col
                if (nums[i][0] > max) max = nums[i][0];
            }
    
            while (pq.Count == nums.Count)
            {
                var cur = pq.Dequeue();
                if (range > max - cur[0])
                {
                    range = max - cur[0];
                    start = cur[0];
                    end = max;
                }
                // cur[2] is col, cur[1] is row
                if (cur[2] + 1 < nums[cur[1]].Count)
                {
                    cur[2]++;
                    cur[0] = nums[cur[1]][cur[2]];
                    pq.Enqueue(cur);
                    if (cur[0] > max) max = cur[0];
                }
            }
    
            return new[] { start, end };
        }
    
        public class PriorityQueue<T>
        {
            private Comparison<T> comparison;
            private List<T> list = new List<T>();
            public PriorityQueue(Comparison<T> comparison)
            {
                this.comparison = comparison;
            }
    
            public int Count
            {
                get
                {
                    return list.Count;
                }
            }
            public T Peek()
            {
                return list[0];
            }
    
            /// <summary>
            /// A little bit hard to implement sift up/down, 
            /// so just maintain a sorted list via using binary search
            /// O(log(N)) for enqueue 
            /// </summary>
            /// <param name="val"></param>
            public void Enqueue(T val)
            {
                var i = list.BinarySearch(val, Comparer<T>.Create(this.comparison));
                i = i >= 0 ? i : ~i;
                list.Insert(i, val);
            }
    
            public T Dequeue()
            {
                var val = list[0];
                list.RemoveAt(0);
                return val;
            }
        }
    }
    

  • 4
    Q

    Same idea with shorter code in Java. One optimization I made is: when the algorithm found the length of range is 0, it should stop immediately. Hope it helps.

    public int[] smallestRange(List<List<Integer>> nums) {
            PriorityQueue<int[]> pq = new PriorityQueue<int[]>((a, b) -> a[0] - b[0]);
            int max = Integer.MIN_VALUE;
            for(int i = 0; i < nums.size(); i++) {
            	max = Math.max(max, nums.get(i).get(0));
            	pq.add(new int[] {nums.get(i).get(0), i, 0});
            }
            int[] result = {pq.peek()[0], max};
            while(result[1] - result[0] != 0 && pq.peek()[2] + 1 < nums.get(pq.peek()[1]).size()) {
            	int[] curr = pq.poll();
            	pq.add(new int[] {nums.get(curr[1]).get(curr[2] + 1), curr[1], curr[2] + 1});
            	max = Math.max(max, nums.get(curr[1]).get(curr[2] + 1));
            	if(max - pq.peek()[0] < result[1] - result[0]) result = new int[] {pq.peek()[0], max};
            }
            return result;
        }
    

  • 7

  • 0
    B

    Very clear BFS2 solution.


  • 0
    C

    @qswawrq I am new to lambda expressions. Can you tell me what this
    PriorityQueue<int[]> pq = new PriorityQueue<int[]>((a, b) -> a[0] - b[0]);
    means?


  • 0
    Q

    @Raktima It means that pq compares its incoming element (a) and old elements (b) for order based on the value of a[0] - b[0]. The value of a[0] - b[0] can be a negative integer, zero, or a positive integer as a is less than, equal to, or greater than b.


  • 0

    Maintain a queue that have exact one element from each array.
    The range is difference of max and min value in the queue in different phase.
    Compare all the ranges to find the minimum range.

    class Solution {
        public int[] smallestRange(List<List<Integer>> nums) {
           
           PriorityQueue<int[]> pq = new PriorityQueue<int[]>((a,b)->{
             return a[0] - b[0];
           });
          
           int max = Integer.MIN_VALUE, range = Integer.MAX_VALUE;
           for(int i = 0; i < nums.size(); i++){
             List<Integer> list = nums.get(i);
             pq.offer(new int[]{list.get(0), i, 0});
             max = Math.max(max, list.get(0));
           }
           
           int start = 0, end = max;
           while(pq.size() == nums.size()){
             int[] node = pq.poll();
             if(max - node[0] < range){
               start = node[0];
               end = max;
               range = end - start;
             }
             
             List<Integer> list = nums.get(node[1]);
             if(node[2] + 1 < list.size()){
               int[] newNode = {list.get(node[2] + 1), node[1], node[2] + 1};
               pq.offer(newNode);
               max = Math.max(max, newNode[0]);
             }
           }
           return new int[]{start, end};
        }
    }

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.