โปรแกรมหา K คู่ที่มีผลรวมมากที่สุดจากสองอาเรย์ใน Python

ได้รับอาเรย์สองตัว nums0, nums1 และ k หา ผลรวมของ k คู่ที่มี sum มากที่สุด (แต่ละคู่เอา 1 ตัวจากแต่ละอาเรย์)

ตัวอย่าง

nums0 = [8, 6, 12], nums1 = [4, 6, 8], k = 2

Output: 38

คู่มากสุด: (12,8)=20, (12,6)=18 → รวม 38

แนวทางแก้ปัญหา

Max Heap + Visited Set (Best-First Search):

1. เรียงทั้งสองอาเรย์
2. เริ่มจากคู่มากสุด: (nums0[-1], nums1[-1]) push เข้า max heap
3. Pop k ครั้ง: แต่ละครั้ง expand ไปคู่ข้างเคียง (i-1,j) และ (i,j-1) ถ้ายังไม่visited

การใช้งาน (Python)

import heapq

class Solution:
    def solve(self, nums0, nums1, k):
        if k > len(nums0) * len(nums1):
            return 0
        
        nums0.sort()
        nums1.sort()
        
        # max heap: store (-sum, i, j)
        heap = [(-(nums0[-1] + nums1[-1]), len(nums0) - 1, len(nums1) - 1)]
        visited = {(len(nums0) - 1, len(nums1) - 1)}
        ans = 0
        
        for _ in range(k):
            neg_sum, i, j = heapq.heappop(heap)
            ans += -neg_sum
            
            # Expand to (i-1, j)
            if i > 0 and (i - 1, j) not in visited:
                visited.add((i - 1, j))
                heapq.heappush(heap, (-(nums0[i - 1] + nums1[j]), i - 1, j))
            
            # Expand to (i, j-1)
            if j > 0 and (i, j - 1) not in visited:
                visited.add((i, j - 1))
                heapq.heappush(heap, (-(nums0[i] + nums1[j - 1]), i, j - 1))
        
        return ans

# ทดสอบ
ob = Solution()
print(ob.solve([8, 6, 12], [4, 6, 8], 2))  # Output: 38

ความซับซ้อน

• Time Complexity: O(k log k)
• Space Complexity: O(k)

ดีกว่า O(N×M) มากเมื่อ k << N×M