Java 6ms solution beats 90%


  • 0
    M
    public class Solution {
    private static class Pair implements Comparable<Pair> {
        public int i1;
        public int i2;
        public int sum;
        
        public Pair(int i1, int i2, int[] nums1, int[] nums2) {
            this.i1 = i1;
            this.i2 = i2;
            this.sum = nums1[i1] + nums2[i2];
        }
        
        public int compareTo(Pair that) {
            return this.sum - that.sum;
        }
    }
    
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int l1 = nums1.length;
        int l2 = nums2.length;
        List<int[]> result = new ArrayList<>();
        PriorityQueue<Pair> q = new PriorityQueue<>();
        
        if(l1 == 0 || l2 == 0) return result;
        q.offer(new Pair(0, 0, nums1, nums2));
        
        while(k-- > 0 && q.size() > 0) {
            Pair top = q.poll();
            result.add(new int[]{nums1[top.i1], nums2[top.i2]});
            if(top.i2 < l2 - 1)
                q.offer(new Pair(top.i1, top.i2 + 1, nums1, nums2));
            if(top.i2 == 0 && top.i1 < l1 - 1)
                q.offer(new Pair(top.i1 + 1, 0, nums1, nums2));
        }
        
        return result;
    }
    

    }

    The trick that speeds it up is to increase the size of the queue only when necessary (that is, when a pair with index(x,0) has been used, then we need to start considering (x+1,0))


Log in to reply
 

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