Find K Pairs with Smallest Sums ใน C++: หาคู่ที่มีผลรวมน้อยที่สุด k คู่

โจทย์

อาร์เรย์เรียงลำดับ A1, A2 หา k คู่ (u,v) u∈A1, v∈A2 ที่มีผลรวมน้อยที่สุด เรียงตามผลรวม

ตัวอย่าง: A1=[1,7,11], A2=[2,4,6], k=3 → [(1,2),(1,4),(1,6)]

แนวคิด: Min-Heap (Priority Queue)

เริ่มจากคู่ (A1[i], A2[0]) ทุก i ใส่ heap เรียงตาม sum ดึงคู่น้อยสุด แล้ว push (A1[i], A2[idx+1]) ต่อ

ขั้นตอน

1. Struct Data{int a, b, idx} + Comparator เรียง a+b น้อยกว่า
2. Push (A1[i], A2[0], 0) ทุก i ลง heap
3. while !heap.empty() && k-- > 0:
   • pop curr เพิ่ม {curr.a, curr.b} ลงผลลัพธ์
   • ถ้า curr.idx + 1 < A2.size() → push (curr.a, A2[curr.idx+1], curr.idx+1)

ตัวอย่างโค้ด C++

#include <bits/stdc++.h>
using namespace std;

struct Data {
    int a, b, idx;
    Data(int a, int b, int idx) : a(a), b(b), idx(idx) {}
};

struct Comparator {
    bool operator()(const Data& x, const Data& y) {
        return x.a + x.b > y.a + y.b; // min-heap
    }
};

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        priority_queue<Data, vector<Data>, Comparator> pq;
        int n = nums1.size(), m = nums2.size();
        if (!n || !m) return {};
        
        for (int i = 0; i < n; i++) pq.push(Data(nums1[i], nums2[0], 0));
        
        vector<vector<int>> ret;
        while (!pq.empty() && k--) {
            Data curr = pq.top(); pq.pop();
            ret.push_back({curr.a, curr.b});
            if (curr.idx + 1 < m) {
                pq.push(Data(curr.a, nums2[curr.idx + 1], curr.idx + 1));
            }
        }
        return ret;
    }
};

int main() {
    Solution ob;
    vector<int> n1 = {1,7,11}, n2 = {2,4,6};
    auto res = ob.kSmallestPairs(n1, n2, 3);
    for (auto& p : res) cout << "[" << p[0] << "," << p[1] << "] ";
    return 0;
}

อินพุต / เอาต์พุต

Input: [1,7,11] [2,4,6] 3

Output: [[1,2],[1,4],[1,6]]

ความซับซ้อน

• Time: O((n + k) log n) - push เริ่มต้น n, แล้ว k รอบ pop+push
• Space: O(n)