priority_queue 的使用 —— 求第 k 小的和

原题再现

其实一想到第 k 小，马上就要想到 priority_queue！结果，我第一版交了个 C 语言版本上去。一开始的思路想通过直接操纵 a、b 两个数组的下标来输出，但是我考虑得太简单了！认真一想发现这个操纵规则我自己也没搞懂。

接下来我依然没考虑到 priority_queue。我使用了 list，试图在双层嵌套循环中每次都 find_if，然后插入，并控制元素个数不超过，最后返回 list::back()。马上就超时了。

所以这个是大教训！然后我问了副驾驶。马上它就生成了一份代码。我稍微改了下：

using namespace std;
#include 
#include 
#include 
#include 
#include 

struct State {
    State(int i, int j, int sum) :i(i), j(j), s(sum) {}
    bool operator>(const State& st) const { return s > st.s; }
    int i, j;
    int s;
};

int main(void)
{
    cin.sync_with_stdio(false);
    cin.tie(nullptr);

    int m, n, k;
    while (cin >> m >> n >> k) {
        vector a(m), b(n);
        for (int& i : a) {
            cin >> i;
        }
        sort(a.begin(), a.end());
        for (int& i : b) {
            cin >> i;
        }
        sort(b.begin(), b.end());

        priority_queue, greater> min_heap;
        for (int i = 0; i < m; i++) {
            min_heap.push(State(i, 0, a[i] + b[0]));
        }

        int x;
        while (true) {
            State s = min_heap.top();
            min_heap.pop();
            x = s.s;
            if (!--k) {
                break;
            }

            if (s.j + 1 < n) {
                s.s -= b[s.j];
                s.s += b[++s.j];
                min_heap.push(s);
            }
        }
        cout << x << '\n';
    }
    return 0;
}

我第一眼看过去觉得太神奇了，复杂度是。首先保证了第一次弹出的是最小的。只需要再插进去。刚插进去的元素如果更小，下次有可能马上就弹出来了。

这个状态的保存我之前也从来没写过。机器人说这个是惯用法。