LeetCode [786] K-th Smallest Prime Fraction

786. K-th Smallest Prime Fraction

Hard

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] = p and 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.

class Solution {
    int n, p=0, q=1;
    int[] A;
    public int[] kthSmallestPrimeFraction(int[] A, int K) {
        this.n = A.length;
        this.A = A;
        double l = 0, r = 1;
        while(true){
            double m = (l+r)/2;
            int cnt = compute(m);
            if(cnt<K){
                l = m;
            }else if(cnt>K){
                r = m;
            }else{
                return new int[]{p, q};
            }
        }
    }

    //return the number of fractions <= x
    //update p, q such that it records the largest p/q<=x
    int compute(double x){
        p = 0;
        int ret = 0;
        int j = 0;
        for(int i=0; i<n; ++i){
            while(j<n && A[i] > A[j]*x) j++;
            //if j<n, j is the first number satisfying A[i]/A[j]<=x
            //ie., A[i]/A[j] is the largest pair satisfying A[i]/A[j]<=x
            ret += n-1-j+1;
            if(j<n && A[i] * q > A[j] * p){//A[i]/A[j]>p/q
                p = A[i];
                q = A[j];
            }
        }
        return ret;
    }
}