C++ 9ms. num of items in priority_queue is kept minimal.


  • 0
    Y

    First, assume there is a 2d Boolean vector a[], where a[i,j] means nums1[i] and nums2[j] has been selected as a small pair. So we want to add its right and lower neighbors to candidate heap, but only if that neighbor's left and upper neighbors are all selected. This will minimize the heap size.

    Second, we can zip the 2d vector into 1d vector by run-length encoding. This saves space.

    Because heap operation is O(log(n)), we still end up with O(k*log(min(m,n,k))) in worse case. In best case, we got O(k).

    class Solution {
        struct Pair {
            int sum;
            int i1;
            int i2;
            Pair (int sum, int i1, int i2) : sum(sum), i1(i1), i2(i2) { }
            bool operator<(const Pair& o) const { return sum > o.sum || 
                sum == o.sum && i1 > o.i1 || sum == o.sum && i1 == o.i1 && i2 > o.i2; }
        };
    
        template <bool ArrayFirst>
        vector<pair<int, int>> pairsWithOneArray(vector<int>& nums, int other, int k) {
            vector<pair<int, int>> result;
            for (size_t i = 0; i < std::min(nums.size(), (size_t)k); ++i) {
                if (ArrayFirst)
                    result.emplace_back(nums[i], other);
                else
                    result.emplace_back(other, nums[i]);
            }
            return result;
        }
    
    public:
        virtual vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
            if (nums1.empty() || nums2.empty() || k <= 0) return vector<pair<int, int>>();
            if (nums1.size() == 1) {
                return pairsWithOneArray<false>(nums2, nums1[0], k);
            }
            if (nums2.size() == 1) {
                return pairsWithOneArray<true>(nums1, nums2[0], k);
            }
    
            vector<pair<int, int>> result;
            result.emplace_back(nums1[0], nums2[0]);
            if (k == 1) return result;
    
            if (nums1.size() > (size_t)k) nums1.resize(k);
            if (nums2.size() > (size_t)k) nums2.resize(k);
    
            priority_queue<Pair> heap;
            heap.emplace(nums1[0] + nums2[1], 0, 1);
            heap.emplace(nums1[1] + nums2[0], 1, 0);
    
            vector<int> nextUnused(nums1.size(), 0);
            nextUnused[0] = 1;
    
            do {
                auto p = heap.top();
                result.emplace_back(nums1[p.i1], nums2[p.i2]);
                if (result.size() == k) return result;
                heap.pop();
                nextUnused[p.i1] = p.i2 + 1;
                if (p.i1 < (int)nums1.size() - 1) {
                    if (nextUnused[p.i1 + 1] == p.i2) {
                        heap.emplace(nums1[p.i1 + 1] + nums2[p.i2], p.i1 + 1, p.i2);
                    }
                }
                if (p.i2 < (int)nums2.size() - 1) {
                    if (!p.i1 || nextUnused[p.i1 - 1] >= p.i2 + 2) {
                        heap.emplace(nums1[p.i1] + nums2[p.i2 + 1], p.i1, p.i2 + 1);
                    }
                }
            } while (!heap.empty());
    
            return result;
        }
    };
    

Log in to reply
 

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