786. K-th Smallest Prime Fraction

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.

思路：priority queue，Python会TLE

import heapq
class Solution:
    def kthSmallestPrimeFraction(self, A, K):
        """
        :type A: List[int]
        :type K: int
        :rtype: List[int]
        """
        q, n = [], len(A)
        for i in range(1, n): heapq.heappush(q, (1.0*A[0]/A[i], 0, i))
        
        while K:
            _,i,j = heapq.heappop(q)
            heapq.heappush(q, (1.0*A[i+1]/A[j], i+1, j)) # i will not be larger than n, there must have a valid solution
            K-=1
            if K==0: return [A[i], A[j]]
        
        return [0,0]

有2点注意：

1. Python heapq无法自定义comparator，只能预先计算出值放到第一个元素

2. 因为都是prime number，所以得到的数没有重复

3. Java 要避免浮点数运算才能AC

import java.util.Comparator;
import java.util.PriorityQueue;

class Solution {
    public int[] kthSmallestPrimeFraction(final int[] A, int K) {
        PriorityQueue<int[]> pq = new PriorityQueue<int[]>(K, new Comparator<int[]>(){
			public int compare(int[] a1, int[] a2) {
//				double d = 1.0*A[a1[0]]/A[a1[1]] - 1.0*A[a2[0]]/A[a2[1]];
//				if(d>0) return 1;
//				else if(d<0) return -1;
//				else return 0;
				return A[a1[0]]*A[a2[1]] - A[a2[0]]*A[a1[1]];
			}
        });
        
        for(int i=1; i<A.length; i++)	pq.add(new int[]{0, i});
        while(K>0) {
        	int[] t = pq.poll();
        	pq.add(new int[]{t[0]+1, t[1]});
        	K--;
        	if(K==0) return new int[]{A[t[0]], A[t[1]]};
        }
        return null;
    }
}