leetcode Find K Pairs with Smallest Sums

leetcode 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:
1
2
3
4
5
6
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]
Example 2:
1
2
3
4
5
6
Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

The first 2 pairs are returned from the sequence:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
Example 3:
1
2
3
4
5
6
Given nums1 = [1,2], nums2 = [3],  k = 3 

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]

题目地址：leetcode Find K Pairs with Smallest Sums

题意：给定两个已排好序的数组，要求从两个数组中分别取出一个数，构成一个对，求所有能组成的对中，最小的K个。

思路：

最直接的思路是把全部的数都组合一下，然后维护一个大小为k的最大堆这样复杂度O(mnlogk)

其实还可以维护一个最小堆，每次取出堆顶的元素，然后把该元素相邻结点加入进去。这样能保证最小。

C++

struct Order {
	int sum;
	int index1, index2;
	Order(int a, int b, int sum) : index1(a), index2(b), sum(sum) {}
	bool operator < (const Order& b) const {
		return sum > b.sum;
	}
};
int dx[] = { 1,0 };
int dy[] = { 0,1 };
class Solution {
public:
	vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
		vector<pair<int, int>> ans;
		if (nums1.empty() || nums2.empty()) return ans;

		int n = nums1.size(), m = nums2.size();
		vector<vector<bool>> vis(n, vector<bool>(m, false));
		priority_queue<Order> q;

		q.push(Order(0, 0, nums1[0] + nums1[1]));
		vis[0][0] = true;

		while (k > 0 && !q.empty()) {
			Order cur = q.top();
			q.pop();
			k--;
			int x = cur.index1, y = cur.index2;
			ans.push_back(make_pair(nums1[x], nums2[y]));
			for (int i = 0; i < 2; i++) {
				int nx = x + dx[i], ny = y + dy[i];
				if (nx < n && ny < m && !vis[nx][ny]) {
					q.push(Order(nx, ny, nums1[nx] + nums2[ny]));
					vis[nx][ny] = true;
				}
			}
		}
		return ans;
	}
};

更好的写法，先将nums1中每个数和nums2的第一个数放进堆，这样，每次从堆中取cur，下一个取cur.i和cur.j+1下标位置的即可

struct Node{
    int sum, i, j;
    Node(int sum, int i, int j):sum(sum), i(i), j(j){}
    bool operator < (const Node &b) const{
        return sum > b.sum;
    }
};

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int>> ans;
        if(nums1.empty()  nums2.empty()) return ans;
        priority_queue<Node> q;
        for(int i = 0; i < k && i < nums1.size(); ++i)
            q.push(Node(nums1[i] + nums2[0], i, 0));

        while(k-- && !q.empty()){
            Node cur = q.top();
            q.pop();
            ans.push_back(make_pair(nums1[cur.i], nums2[cur.j]));
            if(cur.j + 1 < nums2.size())
                q.push(Node(nums1[cur.i] + nums2[cur.j + 1], cur.i, cur.j + 1));
        }
        return ans;
    }
};

本题是leetcode 373 Find K Pairs with Smallest Sums 题解

更多题解可以查看：https://www.hrwhisper.me/leetcode-algorithm-solution/