Klog(2K) java solution


  • 0
    Z

    The idea is when you process the nums1[i] and nums2[j], add the nums[i+1][j] and nums[i][j+1] element to the priorityqueue.

    public class Solution {
        public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
            PriorityQueue<int[]> pq = new PriorityQueue<int[]>(k,(a,b) -> (nums1[a[0]]+nums2[a[1]]-nums1[b[0]]-nums2[b[1]]));
            List<int[]> res = new ArrayList<int[]>();
            if(nums1.length == 0 || nums2.length == 0 || k<= 0) return res;
            int m = nums1.length;
            int n = nums2.length;
            boolean[][] A = new boolean[m][n];
            pq.add(new int[]{0,0});
            A[0][0] = true;
            for(int i=0;i<m*n && i< k;i++) {
                int[] cur = pq.poll();
                res.add(new int[]{nums1[cur[0]], nums2[cur[1]]});
                int a = cur[0];
                int b = cur[1];
                if(a+1 <m && !A[a+1][b]) {
                    A[a+1][b] = true;
                    pq.add(new int[]{a+1,b});
                }
                if(b+1 < n && !A[a][b+1]) {
                    A[a][b+1] = true;
                    pq.add(new int[]{a,b+1});
                }
            }
            return res;
        }
    }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.