[LeetCode] 786. K-th Smallest Prime Fraction 第K小的质分数

A sorted list A contains 1, plus some number of primes. Then, for every p < q in the list, we consider the fraction p/q.

What is the K-th smallest fraction considered? Return your answer as an array of ints, where answer[0] = pand answer[1] = q.

Examples:
Input: A = [1, 2, 3, 5], K = 3
Output: [2, 5]
Explanation:
The fractions to be considered in sorted order are:
1/5, 1/3, 2/5, 1/2, 3/5, 2/3.
The third fraction is 2/5.

Input: A = [1, 7], K = 1
Output: [1, 7]

Note:

A will have length between 2 and 2000.
Each A[i] will be between 1 and 30000.
K will be between 1 and A.length * (A.length - 1) / 2.

这道题给了我们一个有序数组，里面是1和一些质数，说是对于任意两个数，都可以组成一个 [0, 1] 之间分数，让求第K小的分数是什么，题目中给的例子也很好的说明了题意。那么最直接暴力的解法就是遍历出所有的分数，然后再进行排序，返回第K小的即可。但是这种无脑暴力搜索的方法 OJ 是不答应的，无奈，只能想其他的解法。由于数组是有序的，所以最小的分数肯定是由第一个数字和最后一个数字组成的，而接下来第二小的分数就不确定是由第二个数字和最后一个数字组成的，还是由第一个数字跟倒数第二个数字组成的。这里用一个最小堆来存分数，那么每次取的时候就可以将最小的分数取出来，由于前面说了，不能遍历所有的分数都存入最小堆，那么该怎么办呢，可以先存n个，哪n个呢？其实就是数组中的每个数字都和最后一个数字组成的分数。由于需要取出第K小的分数，那么在最小堆中取K个分数就可以了，第一个取出的分数就是那个由第一个数字和最后一个数字组成的最小的分数，然后就是精髓所在了，此时将分母所在的位置前移一位，还是和当前的分子组成新的分数，这里即为第一个数字和倒数第二个数字组成的分数，存入最小堆中，那么由于之前已经将第二个数字和倒数第一个数字组成的分数存入了最小堆，所以不用担心第二小的分数不在堆中，这样每取出一个分数，都新加一个稍稍比取出的大一点的分数，这样取出了第K个分数即为所求，参见代码如下：

解法一：

class Solution {
public:
    vector<int> kthSmallestPrimeFraction(vector<int>& A, int K) {
        priority_queue<pair<double, pair<int, int>>> q;
        for (int i = 0; i < A.size(); ++i) {
            q.push({-1.0 * A[i] / A.back(), {i, A.size() - 1}});
        }
        while (--K) {
            auto t = q.top().second; q.pop();
            --t.second;
            q.push({-1.0 * A[t.first] / A[t.second], {t.first, t.second}});
        }
        return {A[q.top().second.first], A[q.top().second.second]};
    }
};

其实这道题比较经典的解法是用二分搜索法 Binary Search，使用的二分搜索法是博主归纳总结帖 LeetCode Binary Search Summary 二分搜索法小结中的第四种，即二分法的判定条件不是简单的大小关系，而是可以抽离出子函数的情况，下面来看具体怎么弄。这种高级的二分搜索法在求第K小的数的时候经常使用，比如 Kth Smallest Element in a Sorted Matrix，Kth Smallest Number in Multiplication Table，和 Find K-th Smallest Pair Distance 等。思路都是用 mid 当作 candidate，然后统计小于 mid 的个数 cnt，和K进行比较，从而确定折半的方向。这道题也是如此，mid 为候选的分数值，刚开始时是 0.5，然后需要统计出不大于 mid 的分数都个数 cnt，同时也需要找出最接近 mid 的分数，当作返回的候选值，因为一旦 cnt 等于K了，直接将这个候选值返回即可，这个候选值分数是由p和q来表示的，其中p表示分子，初始化为0，q表示分母，初始化为1（因为除数不能为0），在内部的 while 循环退出时，分数 A[i]/A[j] 就是最接近 mid 的候选者，此时假如 p/q 要小于 A[i]/A[j]，就要分别更新p和q。否则如果 cnt 小于K，说明应该增大一些 mid，将 left 赋值为 mid，反之如果 cnt 大于K，需要减小 mid，将 right 赋值为 mid，参见代码如下：

解法二：

class Solution {
public:
    vector<int> kthSmallestPrimeFraction(vector<int>& A, int K) {
        double left = 0, right = 1;
        int p = 0, q = 1, cnt = 0, n = A.size();
        while (true) {
            double mid = left + (right - left) / 2.0;
            cnt = 0; p = 0;
            for (int i = 0, j = 0; i < n; ++i) {
                while (j < n && A[i] > mid * A[j]) ++j;
                cnt += n - j;
                if (j < n && p * A[j] < q * A[i]) {
                    p = A[i];
                    q = A[j];
                }
            }
            if (cnt == K) return {p, q};
            if (cnt < K) left = mid;
            else right = mid;
        }
    }
};