leetcode#786. K-th Smallest Prime Fraction

786. K-th Smallest Prime Fraction

Problem Description

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.

Analysis and Solution

The brute force solution is just list all the possible fractions and sort them. However, there are n*(n-1)/2 elements in total, causes a O(n^3) time complexity to sort them all. So we need to find a faster method to sort.

At first, a lemma is clearly introduced:

When the numerator is fixed, the fraction with larger denominator is always smaller.

According to the lemma, an easy thought is that we divide each possible fraction into n groups with the same numerator and merge them to find the kth elements. We can imagine each group as a chain with a pointer initially pointed to the largest denominator. Then we maintain a priority queue to pick the smallest fraction among the n marked fractions and move the pointer left, until the kth fraction is found. The total time complexity is O(klogn).

Here is the code:

class Solution {
public:
    vector<int> kthSmallestPrimeFraction(vector<int>& A, int K) {
        vector<int> flag(A.size(), A.size()-1);
        priority_queue<Fraction> q;

        // Initialize priority queue
        for (int i=0; i<A.size(); i++) {
            if (i < flag[i]) {
                q.push(Fraction(A[i], A[flag[i]], i));
            }
        }

        int count = 0;
        Fraction result;
        while (count++ != K) {
            result = q.top();
            q.pop();
            int i = result.count;
            if (i < --flag[i]) {
                q.push(Fraction(A[i], A[flag[i]], i));
            }
        }

        vector<int> tmp;
        tmp.push_back(result.up);
        tmp.push_back(result.down);

        return tmp;
    }
private:
    struct Fraction {
        int up;
        int down;
        int count;
        Fraction (int a=0, int b=0, int c=0) : up(a), down(b), count(c) {};

        friend bool operator<(Fraction x, Fraction y) {
            float xf = float(x.down) / float(x.up);
            float yf = float(y.down) / float(y.up);
            return xf < yf;
        }
    };
};