Tuesday, April 11, 2017

Kth Smallest Sum In Two Sorted Arrays

这道题特别有意思
本质上就是Kth Smallest Element in a Sorted Matrix
但是很难想到
做法都是一样的
写了一个bug就是在加入a+1和b+1的时候要提前check有没有越界,越界了就不能再加入了

public class Solution {
    /**
     * @param A an integer arrays sorted in ascending order
     * @param B an integer arrays sorted in ascending order
     * @param k an integer
     * @return an integer
     */
    private class Tuple {
        public int a, b;
        public int sum;
        public Tuple(int a, int b, int sum) {
            this.a = a;
            this.b = b;
            this.sum = sum;
        }
    }
    public int kthSmallestSum(int[] A, int[] B, int k) {
        // Write your code here
        if (A == null || A.length == 0 || B == null || B.length == 0 || k <= 0) throw new IllegalArgumentException();
        //min heap
        PriorityQueue<Tuple> pq = new PriorityQueue<>(k, new Comparator<Tuple>() {
            public int compare(Tuple t1, Tuple t2) {
                return t1.sum - t2.sum;
            }
        });
        pq.add(new Tuple(0, 0, A[0] + B[0]));
        int count = k;
        Tuple curr = null;
        while (count > 0) {
            curr = pq.poll();
            count--;
            //System.out.println("a:" + curr.a + " b:" + curr.b + " count:" + count);
            if (curr.a == 0) {
                if(curr.b + 1 < B.length) pq.offer(new Tuple(curr.a, curr.b + 1, A[curr.a] + B[curr.b + 1]));
            }
            if (curr.a + 1 < A.length) pq.offer(new Tuple(curr.a + 1, curr.b, A[curr.a + 1] + B[curr.b]));
        }
        return curr.sum;
       
    }
}

No comments:

Post a Comment