373. Find K Pairs with Smallest Sums

class Solution:
    def kSmallestPairs(self, A: List[int], B: List[int], k: int) -> List[List[int]]:
        """tc O(KlgK)  sc O(K)
        main idea: min PQ,   (pair_sum,i_A,j_B) 
        constrains: (1) for each array, we consider at most k elements 
                    (2) for each val in A, if its smallest pairt in B at pos j is pop off, its next smallest pair is at pos j+1 
                    (3) in the beginning, use smallest val in B to get k pairs in A 
        1. take smallest in A as start, get k pairs in B, push into PQ 
        2. each time pop one pair, update res, update cnt, increment index in B (if within boundry), push bach to pq 
        3. stop when eighter count reaches k or pq is empty already 
        """
        res = [] 
        pq = [] # (pair_sum, i_A, j_B)
        j = 0 
        max_level = min(k,len(A)) # within valid range, put into heap no more than k  
        for i in range(max_level):
            pair_sum = A[i] + B[j]
            heappush(pq,(pair_sum,i,j))
        
        cnt = 0 
        while cnt < k and len(pq)>0: # not assume there will be an answer (maybe number of pairs  < k )  [1,2] [3] k =3 
            pair_sum,i,j = heappop(pq)
            res.append([A[i],B[j]])
            cnt += 1 
            if j+1 == len(B):
                continue 
            j += 1 
            heappush(pq,(A[i]+B[j],i,j))
            
        return res