多路归并的一些题目总结

 

写在前面

参考了:

【多路归并】从朴素优先队列到多路归并;

基本问题

88. 合并两个有序数组;

双指针归并:

class Solution {
public:
    void merge(vector<int>& nums1, int m, vector<int>& nums2, int n) {
        // 反着遍历插入数组
        for (int i{m + n - 1}, j{m - 1}, k{n - 1}; i >= 0; --i) {
            if (j < 0)
                nums1[i] = nums2[k--];
            else if (k < 0)
                nums1[i] = nums1[j--];
            else if (nums1[j] > nums2[k])
                nums1[i] = nums1[j--];
            else
                nums1[i] = nums2[k--];
        }
    }
};

264. 丑数 II;

直接队列就显得很方便, 但是速度不行:

class Solution {
public:
    int nthUglyNumber(int n) {
        priority_queue<long, vector<long>, greater<>> pq;
        unordered_set<long> st;
        pq.emplace(1);
        long ans{};
        for (int i{}; i < n; ++i) {
            ans = pq.top();
            pq.pop();
            for (int factor : {2, 3, 5}) {
                auto tmp = ans * factor;
                if (!st.count(tmp)) 
                    st.insert(tmp), pq.emplace(tmp);
            }
        }
        return ans;
    }
};

重点要说的就是多路归并算法: (感觉本质上就是双指针的推广, 多指针往后挪)

class Solution {
public:
    int nthUglyNumber(int n) {
        // 多路归并
        long ans[n];
        memset(ans, 0, sizeof(ans));
        ans[0] = 1;
        for (int i2{}, i3{}, i5{}, i{1}; i < n; ++i) {
            auto a{ans[i2] * 2}, b{ans[i3] * 3}, c{ans[i5] * 5};
            auto x = min(a, min(b, c));
            if (x == a) ++i2;
            if (x == b) ++i3;
            if (x == c) ++i5; // 去重
            ans[i] = x;
        }
        return ans[n - 1];
    }
};

进阶问题

373. 查找和最小的 K 对数字;

这个用纯优先队列会超时, 只能上多路归并了(二分也可以, 看起来很复杂)

class Solution { // 超时的代码
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2,
                                       int k) {
        using pivi = pair<int, vector<int>>;
        priority_queue<pivi, vector<pivi>, greater<>> pq;
        for (auto i : nums1)
            for (auto j : nums2) // 时间复杂度主要来自这里
                pq.emplace(i + j, vector<int>{i, j});
        vector<vector<int>> ans;
        ans.reserve(k);
        for (int i{}; !pq.empty() && i < k; pq.pop(), ++i)
            ans.emplace_back(pq.top().second);
        return ans;
    }
};

下面是优化过的多路归并+优先队列:

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2,
                                       int k) {
        bool flg{true};
        vector<vector<int>> ans;
        ans.reserve(k); // 减少 vector 内存分配开销
        int m = nums1.size(), n = nums2.size();
        if (m > n) swap(nums1, nums2), swap(m, n), flg = false;
        auto cmp = [&](const auto& a, const auto& b) {
            return nums1[a.first] + nums2[a.second] >
                   nums1[b.first] + nums2[b.second];
        };
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)>
            pq(cmp);
        for (int i{}; i < min(m, k); ++i) pq.emplace(i, 0);
        for (; ans.size() < k && !pq.empty(); pq.pop()) {
            auto [a, b] = pq.top();
            ans.emplace_back(flg ? vector<int>{nums1[a], nums2[b]}
                                 : vector<int>{nums2[b], nums1[a]});
            if (b + 1 < n) pq.emplace(a, b + 1);
        }
        return ans;
    }
};

786. 第 K 个最小的素数分数;(其实也是堆来操作的)

class Solution {
public:
    vector<int> kthSmallestPrimeFraction(vector<int>& arr, int k) {
        int n = arr.size();
        using ptype = pair<int, int>;
        auto cmp = [&](const auto& lhs, const auto& rhs) {
            return arr[lhs.first] * arr[rhs.second] >
                   arr[lhs.second] * arr[rhs.first];
        };
        priority_queue<ptype, vector<ptype>, decltype(cmp)> pq(cmp);

        for (int i{1}; i < n; ++i) 
            pq.emplace(0, i);
        for (int _{1}; _ < k; ++_) {
            auto [i, j] = pq.top();
            pq.pop();
            if (i + 1 < j) 
                pq.emplace(i + 1, j);
        }
        return {arr[pq.top().first], arr[pq.top().second]};
    }
};

719. 找出第 K 小的数对距离;

1439. 有序矩阵中的第 k 个最小数组和;