373. Find K Pairs with Smallest Sums

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

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

Find the k pairs (u1,v1),(u2,v2) …(uk,vk) with the smallest sums.

Example 1:

Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

这是展示如何从最简单的方法一步步改进的例子。

首先明白一个基本问题,要找k个sum最小的pair,用heap,并且重新定义heap的compare。用java8的写法:

PriorityQueue que = new PriorityQueue((a, b) -> a[2] - b[2])

其中a[2]-b[2]这是根据把什么样的数据结构存入heap来的,不管用什么样的数据结构,让heap比较两个pair的和

接下来首先尝试brute force:

        List<int[]> ans = new ArrayList<int[]>();
        PriorityQueue<int[]> que = new PriorityQueue<>((a,b) -> a[0] + a[1] - b[0] - b[1]);
        for (int n1 : nums1) {
            for (int n2 : nums2) {
                que.add(new int[]{n1, n2}); //在heap中存入(nums1[i], nums2[j])
            }
        }

        while (k-- > 0 && !que.isEmpty()) {
            ans.add(que.poll());
        }
        return ans;

因为两个array都是sorted,而结果只需要头k个sum最小的pair,一种改进是:先存入nums1[i] + nums2[0], i = 0,…,nums1.length-1. 对每个nums1[i], sum最小的pair除了当前已经存入(i,j)的pair之外,唯一的选择就是nums1[i] + nums2[j+1]了,将这个选择存入heap,同时从heap中取出k个pair

        List<int[]> ans = new ArrayList<int[]>();
        PriorityQueue<int[]> que = new PriorityQueue<>((a,b) -> a[0] + a[1] - b[0] - b[1]);
        for (int n1 : nums1) {
            que.add(new int[]{n1, nums2[0], 0});
        }

        while (k-- > 0 && !que.isEmpty()) {
            int[] cur = que.poll();
            ans.add(new int[]{cur[0], cur[1]});
            if (cur[2] == nums2.length-1) continue;
            que.add(new int[]{cur[0], nums2[cur[2]+1], cur[2]+1}); //对nums1[cur[0]]来说,最小的除了nums2[cur[2]]外就是nums2[cur[2]+1]了
        }
        return ans;

不难发现其实不需要在一上来就在heap中存所有nums1[i]和nums2[0]的pair,只需要加入k个就够了,后面的都不可能在答案中。简单改进:

        List<int[]> ans = new ArrayList<int[]>();
        PriorityQueue<int[]> que = new PriorityQueue<>((a,b) -> a[0] + a[1] - b[0] - b[1]);
        for (int i = 0; i < k && i < nums1.length; i++) {             que.add(new int[]{nums1[i], nums2[0], 0});          }                  while (k-- > 0 && !que.isEmpty()) {
            int[] cur = que.poll();
            ans.add(new int[]{cur[0], cur[1]});
            if (cur[2] == nums2.length-1) continue;
            que.add(new int[]{cur[0], nums2[cur[2]+1], cur[2]+1});
        }
        return ans;

最后一步改进比较难想。首先有一个发现:两个array每个元素相加构成了一个matrix
1 7 11
2| 3 9 13
4| 5 11 15
6| 7 13 17
按小到大排列是“/”型的:3 -> (5, 9) -> (7, 11, 13) -> (13, 15) -> 17
但按这个顺序遍历到所有,或是前k个元素?
想到了tree:
3 -> 9 -> 13

5 ->11 -> 15

7 ->13 -> 17
那么按之前的顺序遍历就成了BFS

        List<int[]> ans = new ArrayList<int[]>();
        if (nums1.length == 0 || nums2.length == 0 || k == 0) return ans;
        PriorityQueue<int[]> que = new PriorityQueue<>((a,b) -> a[2] - b[2]);
        que.add(new int[]{0, 0, nums1[0] + nums2[0]}); // i, j, sum
        while (k-- > 0 && !que.isEmpty()) {
            int[] cur = que.poll();
            int i = cur[0];
            int j = cur[1];
            ans.add(new int[]{nums1[i], nums2[j]});
            if (j < nums2.length - 1) {
                que.add(new int[]{i, j + 1, nums1[i] + nums2[j+1]});
            }
            if (j == 0 && i < nums1.length - 1) {
                que.add(new int[]{i + 1, j, nums1[i+1] + nums2[j]});
            } //上面tree的结构决定的

        }
        return ans;
Advertisements

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s