# 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.
import heapq


class Solution(object):
    def kthSmallestPrimeFraction(self, A, K):
        """
        :type A: List[int]
        :type K: int
        :rtype: List[int]
        """
        n, heap = len(A), []
        for k, v in enumerate(A):
            for i in xrange(n - 1, k, -1):
                t = -A[k] * 1.0 / A[i]
                if len(heap) < K or t < heap[-1]:
                    heapq.heappush(heap, (t, A[k], A[i]))

                if len(heap) > K:
                    heapq.heappop(heap)

        a, b, c = heapq.heappop(heap)

        return [b, c]

    # time : O(NlogN + klogN), initialize pq : O(NlogN), for loop : O(KlogN)
    # space: O(N)
    def kthSmallestPrimeFraction2(self, A, K):
        """
        :type A: List[int]
        :type K: int
        :rtype: List[int]
        """
        # firstly put the smallest number candidates in pq
        pq = [(A[0] * 1.0 / A[i], 0, i) for i in xrange(len(A) - 1, 0, -1)]

        for _ in xrange(K - 1):
            frac, a, b = heapq.heappop(pq)
            a += 1
            if a < b:
                heapq.heappush(pq, (A[a] * 1.0 / A[b], a, b))

        return A[pq[0][1]], A[pq[0][2]]



if __name__ == '__main__':
    print Solution().kthSmallestPrimeFraction([1, 2, 3, 5], 3)
    print Solution().kthSmallestPrimeFraction2([1, 2, 3, 5], 3)
    print Solution().kthSmallestPrimeFraction([1, 2, 3, 5], 4)
    print Solution().kthSmallestPrimeFraction2([1, 2, 3, 5], 4)
    print Solution().kthSmallestPrimeFraction([1, 2, 3, 5], 5)
    print Solution().kthSmallestPrimeFraction2([1, 2, 3, 5], 5)
    print Solution().kthSmallestPrimeFraction([1, 2, 3, 5], 6)
    print Solution().kthSmallestPrimeFraction2([1, 2, 3, 5], 6)
    print Solution().kthSmallestPrimeFraction([1, 7], 1)
    print Solution().kthSmallestPrimeFraction2([1, 7], 1)
