1. 程式人生 > >373. Find K Pairs with Smallest Sums

373. Find K Pairs with Smallest Sums

public class Solution {
	class Pair{
		int[] pair;
		int idx; //numArry2[curIndex]
		long sum;
		public Pair(int idx, int n1, int n2) {
			super();
			this.idx = idx; //note current searched index of array2
			pair = new int[]{n1, n2};
			sum = (long)n1 + n2;
		}
	}
	class ComPair implements Comparator{
		public int compare(Pair o1, Pair o2) {
			return Long.compare(o1.sum, o2.sum); //o1.sum - o2.sum
		}
	}
	public List kSmallestPairs(int[] nums1, int[] nums2, int k){
		List ret = new ArrayList<>();
		if(nums1 == null || nums2 == null || nums1.length == 0 || nums2.length == 0){
			return ret;
		}
		PriorityQueue q = new PriorityQueue<>(k, new ComPair());
		for(int i = 0; i < nums1.length && i < k; i++){
			q.offer(new Pair(0, nums1[i], nums2[0])); //klogk
		}
		for(int i = 0; i < k && !q.isEmpty(); i++){
			Pair p = q.poll(); //remove head of the queue O(logk)
			ret.add(p.pair);
			if(p.idx < nums2.length-1){
				q.offer(new Pair(p.idx+1, p.pair[0], nums2[p.idx+1]));
			}
		}
		return ret;
	}
}