多个有序数组中找第k小的数

用最小堆保存行索引,用colIdx[]保存各行的列索引

int kthSmallest(vector<vector<int>>& matrix, int k) {
    const int N = matrix.size();
    vector<int> colIdx(N, 0);
    auto cmp = [&](int r1, int r2) { // 哪一行的当前元素较小
        if (colIdx[r1] == N) return true; // 相当于r1行元素无穷大
        if (colIdx[r2] == N) return false;
        return matrix[r1][colIdx[r1]] > matrix[r2][colIdx[r2]];
    };
    priority_queue<int, vector<int>, decltype(cmp)> pq(cmp);
    for (int i = 0; i < N; i++) {
        pq.push(i);
    }

    while (true) {
        int minR = pq.top();
        pq.pop();
        k--;
        if (k == 0) return matrix[minR][colIdx[minR]];

        colIdx[minR]++;
        pq.push(minR);
    }
}

注:这题也可用二分搜索解。

两个有序数组的前k个最小和对

把两有序数组一个视作行指示、一个视作列指示,和值作为矩阵元素,就变成行列分别有序的矩阵,求其中前k个最小数。

vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
    if (nums1.empty() || nums2.empty()) return {};
    // nums1的数作行、nums2的数作列、两数和作为矩阵元素,构成行列分别有序的矩阵
    // 用最小堆保存行列索引
    auto cmp = [&](pair<int, int> &p1, pair<int, int> &p2) {
        return nums1[p1.first] + nums2[p1.second] > nums1[p2.first] + nums2[p2.second];
    };
    priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> pq(cmp);

    const int M = nums1.size(), N = nums2.size();
    for (int i = 0; i < M; i++) pq.push({ i, 0 }); // 第一列先入队

    vector<pair<int, int>> ans;
    while (!pq.empty() && k--) {
        auto p = pq.top(); pq.pop();
        ans.push_back({ nums1[p.first], nums2[p.second] });
        if (p.second + 1 < N) pq.push({ p.first, p.second + 1 });
    }
    return ans;
}

至少包含各有序数组一个数的最小范围

vector<int> smallestRange(vector<vector<int>>& nums) {
    const int N = nums.size();
    vector<int> idx(N, 0); // idx[]存各数组的当前下标,数组i的当前元素是nums[i][idx[i]]
    auto cmp = [&](int i, int j) { // 最小堆
        return nums[i][idx[i]] > nums[j][idx[j]];
    };
    priority_queue<int, vector<int>, decltype(cmp)> pq(cmp);
    int rangeEnd = INT_MIN;
    for (int i = 0; i < N; i++) {
        if (nums[i].empty()) break;
        pq.push(i);
        rangeEnd = max(rangeEnd, nums[i][0]);
    }
    if (pq.size() < N) return { INT_MIN, INT_MAX };

    vector<int> ans;
    int minRangeSize = INT_MAX;
    while (true) {
        int minQ = pq.top(); pq.pop();
        int rangeBegin = nums[minQ][idx[minQ]];
        int rangeSize = rangeEnd - rangeBegin + 1;
        if (rangeSize < minRangeSize) {
            minRangeSize = rangeSize;
            ans = { rangeBegin, rangeEnd };
        }
        // 移动指针,下一元素
        idx[minQ]++;
        if (idx[minQ] == nums[minQ].size()) break;
        pq.push(minQ);
        rangeEnd = max(rangeEnd, nums[minQ][idx[minQ]]);
    }
    return ans;
}
大数组中包含小数组所有元素的最小范围

为小数组中每个元素记录大数组中的位置链表num=>idxList,变成求至少包含各有序链表一个数的最小范围。见ctci p589。