LeetCode 786: K-th Smallest Prime Fraction

Problem Restatement

We are given a sorted array arr.

The array contains 1 and some prime numbers. All values are unique.

For every pair of indices i and j where:

0 <= i < j < len(arr)

we form the fraction:

arr[i] / arr[j]

Return the k-th smallest fraction.

The answer should be returned as:

[numerator, denominator]

Input and Output

Function shape:

class Solution:
    def kthSmallestPrimeFraction(self, arr: list[int], k: int) -> list[int]:
        ...

Examples

Example 1:

arr = [1, 2, 3, 5]
k = 3

All valid fractions are:

1/2, 1/3, 1/5, 2/3, 2/5, 3/5

Sorted by value:

1/5, 1/3, 2/5, 1/2, 3/5, 2/3

The third smallest fraction is:

2/5

So the answer is:

[2, 5]

Example 2:

arr = [1, 7]
k = 1

There is only one fraction:

1/7

So the answer is:

[1, 7]

First Thought: Generate All Fractions

A direct solution is:

1. Generate every fraction arr[i] / arr[j] where i < j.
2. Sort all fractions by value.
3. Return the k-th one.

Code idea:

fractions = []

for i in range(n):
    for j in range(i + 1, n):
        fractions.append((arr[i] / arr[j], arr[i], arr[j]))

fractions.sort()
return [fractions[k - 1][1], fractions[k - 1][2]]

This is simple, but it stores all possible fractions.

Problem With Generating All Fractions

There are:

n * (n - 1) / 2

valid fractions.

So generating and sorting all of them costs:

We can avoid storing every fraction at once.

Key Insight

The array is sorted.

For a fixed denominator arr[j], the fractions are ordered by numerator:

arr[0] / arr[j] < arr[1] / arr[j] < ... < arr[j - 1] / arr[j]

So each denominator gives us a sorted list of fractions.

For example, with:

arr = [1, 2, 3, 5]

The chains are:

denominator 2: 1/2
denominator 3: 1/3, 2/3
denominator 5: 1/5, 2/5, 3/5

We need the k-th smallest value across these sorted chains.

This is like merging multiple sorted lists.

Use a min-heap.

Algorithm

Initialize a heap with the smallest fraction for each denominator.

For every denominator index j from 1 to n - 1, the smallest numerator is index 0.

So push:

arr[0] / arr[j], 0, j

Then repeat k - 1 times:

1. Pop the smallest fraction (value, i, j).
2. Move to the next larger fraction with the same denominator:

   arr[i + 1] / arr[j]
3. Push it if i + 1 < j.

After popping k - 1 fractions, the top of the heap is the k-th smallest fraction.

Correctness

For each denominator arr[j], the fractions using that denominator form a sorted chain:

arr[0]/arr[j], arr[1]/arr[j], ..., arr[j - 1]/arr[j]

The algorithm initially inserts the first and smallest fraction from each chain into the heap.

At every step, the heap contains the smallest unremoved fraction from each chain that still has remaining fractions. Therefore, the smallest fraction in the heap is the smallest unremoved fraction overall.

When the algorithm pops a fraction arr[i] / arr[j], it inserts the next fraction from the same chain, arr[i + 1] / arr[j], if that fraction exists. This preserves the heap invariant.

After removing the smallest fraction k - 1 times, exactly k - 1 smaller fractions have been removed. The smallest remaining fraction is therefore the k-th smallest fraction.

So returning the fraction at the top of the heap is correct.

Complexity

Implementation

import heapq

class Solution:
    def kthSmallestPrimeFraction(self, arr: list[int], k: int) -> list[int]:
        n = len(arr)
        heap = []

        for j in range(1, n):
            heapq.heappush(heap, (arr[0] / arr[j], 0, j))

        for _ in range(k - 1):
            _, i, j = heapq.heappop(heap)

            if i + 1 < j:
                heapq.heappush(heap, (arr[i + 1] / arr[j], i + 1, j))

        _, i, j = heap[0]
        return [arr[i], arr[j]]

Code Explanation

We create a min-heap:

heap = []

For each possible denominator, we insert the smallest fraction using that denominator:

for j in range(1, n):
    heapq.heappush(heap, (arr[0] / arr[j], 0, j))

The tuple stores:

Then we remove the smallest fraction k - 1 times:

for _ in range(k - 1):
    _, i, j = heapq.heappop(heap)

After popping arr[i] / arr[j], the next candidate with the same denominator is:

arr[i + 1] / arr[j]

This is valid only if the numerator index remains smaller than the denominator index:

if i + 1 < j:

So we push it into the heap:

heapq.heappush(heap, (arr[i + 1] / arr[j], i + 1, j))

After k - 1 removals, the heap top is the answer:

_, i, j = heap[0]
return [arr[i], arr[j]]

Avoiding Floating-Point Comparison

The heap version above uses floating-point values. It is accepted in practice for this problem, but we can also avoid floating-point comparison by defining a custom fraction object.

For Python, the floating-point version is usually simpler and clear. Since array values are small enough for LeetCode constraints, it works reliably.

A binary-search solution can also avoid storing many heap entries, but it is more complex. The heap approach is the most direct teaching solution.

Testing

def run_tests():
    s = Solution()

    assert s.kthSmallestPrimeFraction([1, 2, 3, 5], 1) == [1, 5]
    assert s.kthSmallestPrimeFraction([1, 2, 3, 5], 2) == [1, 3]
    assert s.kthSmallestPrimeFraction([1, 2, 3, 5], 3) == [2, 5]
    assert s.kthSmallestPrimeFraction([1, 2, 3, 5], 4) == [1, 2]

    assert s.kthSmallestPrimeFraction([1, 7], 1) == [1, 7]

    assert s.kthSmallestPrimeFraction([1, 13, 17, 59], 6) == [17, 59]

    print("all tests passed")

run_tests()

Test coverage: