หาผลรวมที่เล็กที่สุดอันดับ K จากเมทริกซ์แถวเรียงลำดับใน C++

ปัญหา

กำหนดเมทริกซ์ mat ขนาด m × n ที่แต่ละแถวเรียงลำดับจากน้อยไปมาก (Non-decreasing) และจำนวนเต็ม k ให้เลือก عنصرละ 1 ตัวจากแต่ละแถว เพื่อสร้างอาร์เรย์ แล้วหา ผลรวมที่เล็กที่สุดอันดับ K จากทุกอาร์เรย์ที่เป็นไปได้

ตัวอย่าง: mat = [[1,3,11],[2,4,6]], k = 5

ผลรวมที่เป็นไปได้ (เรียงจากน้อยไปมาก):

• [1,2] → 3
• [1,4] → 5
• [3,2] → 5
• [3,4] → 7
• [1,6] → 7 ← อันดับ 5

คำตอบ: 7

แนวคิดการแก้: Priority Queue + State Expansion

เริ่มจากผลรวมน้อยที่สุด (เลือกตัวแรกของทุกแถว) แล้วขยาย State โดยเพิ่มดัชนีในแถวทีละแถว ใช้ Min-Heap (Priority Queue) ดึงผลรวมน้อยสุดทีละอัน จนครบ K รอบ

State Representation

• เก็บ [sum, idx_0, idx_1, ..., idx_{m-1}] โดย idx_i = ดัชนีที่เลือกในแถว i
• เริ่มต้น: [sum_of_first_elements, 0, 0, ..., 0]
• จาก State ปัจจุบัน สร้าง State ใหม่โดยเพิ่ม idx_i ทีละ 1 (ถ้ายังไม่เกินขอบเขต)

ป้องกัน State ซ้ำ

ใช้ set<vector<int>> เก็บ State ที่เคยดึงออกจาก PQ แล้ว เพื่อไม่ให้ประมวลผลซ้ำ

โค้ด C++

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

struct Cmp {
    bool operator()(vector<int>& a, vector<int>& b) {
        return !(a[0] < b[0]);  // Min-heap by sum
    }
};

class Solution {
    priority_queue<vector<int>, vector<vector<int>>, Cmp> pq;
    vector<vector<int>> m;
    int z;  // จำนวนคอลัมน์
    
    void update(vector<int>& v, int i, bool ok = false) {
        if (i == v.size()) {
            if (!ok) return;
            int sum = 0;
            for (int j = 0; j < v.size(); j++) sum += m[j][v[j]];
            vector<int> temp(v.begin(), v.end());
            temp.insert(temp.begin(), sum);
            pq.push(temp);
            return;
        }
        v[i]++;
        if (!ok && v[i] < z) update(v, i + 1, true);
        v[i]--;
        update(v, i + 1, ok);
    }
    
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        this->m = mat;
        int n = m.size();
        z = m[0].size();
        
        int ret = 0;
        for (int i = 0; i < n; i++) ret += m[i][0];
        
        vector<int> temp(n);
        temp.insert(temp.begin(), ret);
        pq.push(temp);
        
        set<vector<int>> visited;
        
        while (k--) {
            vector<int> cur = pq.top(); pq.pop();
            visited.insert(cur);
            ret = cur[0];
            cur.erase(cur.begin());  // ลบ sum ออก เหลือแค่ indices
            update(cur, 0);
            
            while (!pq.empty() && visited.count(pq.top())) {
                pq.pop();
            }
        }
        return ret;
    }
};

int main() {
    Solution ob;
    vector<vector<int>> v = {{1,3,11},{2,4,6}};
    cout << ob.kthSmallest(v, 5);  // 7
    return 0;
}

ความซับซ้อน

• Time: O(k × m × log(k)) ประมาณ (m = จำนวนแถว)
• Space: O(k × m) สำหรับ PQ และ Visited Set

Input / Output

Input:  {{1,3,11},{2,4,6}}, k=5
Output: 7