leetcode Find K Pairs with Smallest Sums solutions comparisons

cpp

#include <iostream>
#include <algorithm>
#include <queue>
#include <vector>
#include <chrono>
#include <tuple>

using namespace std;

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        // Sanity check
        if (nums1.empty() || nums2.empty())
            return {};

        // Find the Kth sum (will be tracked by "left" variable)
        // Same algorithm as for "smallest element in sorted matrix" problem
        // https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/description/
        int left = nums1.front() + nums2.front(), right = nums1.back() + nums2.back() + 1;
        while (right - 1 > left)
        {
            int m = left + (right - left) / 2, count = 0;
            for (auto e : nums1)
            {
                auto it = std::lower_bound(nums2.begin(), nums2.end(), m - e);
                count += it - nums2.begin();
            }
            if (count < k + 1)
                left = m;
            else
                right = m;
        }

        // First get all pairs with sum < Kth sum
        std::vector<std::pair<int, int>> pairs;
        for (auto e1 : nums1)
            for (auto e2 : nums2)
            {
                if (e1 + e2 >= left)
                    break;
                pairs.emplace_back(e1, e2);
            }
        if (pairs.size() == k)
            return pairs;

        // If there is room left, get as many pairs with sum == Kth sum, as necessary
        for (auto e1 : nums1)
            for (auto e2 : nums2)
            {
                if (e1 + e2 > left)
                    break;
                if (e1 + e2 == left)
                {
                    pairs.emplace_back(e1, e2);
                    if (pairs.size() == k)
                        return pairs;
                }
            }
        return pairs;
    }
};

class Solution2 {

public:

    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {

        vector<pair<int, int>> pairs;

        if (nums1.size() > nums2.size()) {

            vector<pair<int, int>> tmp = kSmallestPairs(nums2, nums1, k);

            for (const auto& pair : tmp) {

                pairs.emplace_back(pair.second, pair.first);

            }

            return pairs;

        }



        using P = pair<int, pair<int, int>>;

        priority_queue<P, vector<P>, greater<P>> q;

        auto push = [&nums1, &nums2, &q](int i, int j) {

            if (i < nums1.size() && j < nums2.size()) {

                q.emplace(nums1[i] + nums2[j], make_pair(i, j));

            }

        };



        push(0, 0);

        while (!q.empty() && pairs.size() < k) {

            auto tmp = q.top(); q.pop();

            int i, j;

            tie(i, j) = tmp.second;

            pairs.emplace_back(nums1[i], nums2[j]);

            push(i, j + 1);

            if (j == 0) {

                push(i + 1, 0);  // at most queue min(m, n) space.

            }

        }

        return pairs; 

    }

};

template <class Solution>
long long Test(std::vector<int> &data1, std::vector<int> &data2, int K)
{
    auto t1 = std::chrono::high_resolution_clock::now();
    auto r = Solution().kSmallestPairs(data1, data2, K);
    auto t2 = std::chrono::high_resolution_clock::now();
    std::cout << "Solution done in: " << std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1).count() << " milliseconds\n";
    long long sum = 0;
    for (auto &k : r)
        sum += k.first + k.second;
    return sum;
}

int main()
{
    int N = 10000, M = 10000, K = N * M;
    std::vector<int> data1(N), data2(M);
    std::iota(data1.begin(), data1.end(), 1);
    std::iota(data2.begin(), data2.end(), 1);
    auto s1 = Test<Solution>(data1, data2, K);
    auto s2 = Test<Solution2>(data1, data2, K);
    if (s1 != s2)
        std::cout << "Different\n";
}