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 between2
and2000
.- Each
A[i]
will be between1
and30000
. K
will be between1
andA.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;
}
}