Problem

Given two integer arrays sorted in ascending order and an integer k. Define sum = a + b, where a is an element from the first array and b is an element from the second one. Find the kth smallest sum out of all possible sums.

Example

Given[1, 7, 11]and[2, 4, 6].

For k =3, return7.

For k =4, return9.

For k =8, return15

思路

• 有时候怀疑自己想问题的逻辑. 会陷入一种一定要用heap来解决问题的思路
• 正确的想法应该是:
  • 跑一跑上面的例子, 如何找到第一个最小的, 然后如何找到第二个最小, 第三个最小, 有什么联系?
  • 因为是两个数组的和, 所以可以画成一个矩阵,然后再分析.
  • 现将问题展开到最大面, 然后再解决

    /*
        3   5   7
        9   11  13
        13  15  17
    */

    class Element {
        int x;
        int y;
        int val;
        public Element (int x, int y, int val) {
            this.x = x;
            this.y = y;
            this.val = val;
        }
    }
    public int kthSmallestSum(int[] A, int[] B, int k) {
        if (A == null || B == null) {
            return -1;
        }

        TreeSet<Element> set = new TreeSet<>(new Comparator<Element>() {
            public int compare (Element e1, Element e2) {
                if (e1.val != e2.val) return e1.val - e2.val;
                if (e1.x != e2.x) return e1.x - e2.x;
                return e1.y - e2.y;
            }
        });

        set.add(new Element(0, 0, A[0] + B[0]));

        while (!set.isEmpty()) {
            Element curt = set.pollFirst();
            k--;
            if (k == 0) {
                return curt.val;
            }
            int x = curt.x + 1;
            int y = curt.y + 1;
            if (x < A.length) {
                Element n = new Element(x, curt.y, A[x] + B[curt.y]);
                set.add(n);
            }
            if (y < B.length) {
                Element n = new Element(curt.x, y, A[curt.x] + B[y]);
                set.add(n);
            }
        }
        return -1;
    }