LeetCode 786 - K-th Smallest Prime Fraction

The problem asks us to find the k-th smallest fraction that can be formed by dividing one element of a sorted array arr by another element later in the array. The array arr is strictly increasing, starts with 1, and all elements after the first are prime numbers.

leetcodemediumarraytwo-pointersbinary-searchsortingheap-(priority-queue)

LeetCode Problem 786

Difficulty: 🟡 Medium
Topics: Array, Two Pointers, Binary Search, Sorting, Heap (Priority Queue)

Solution

Problem Understanding

The problem asks us to find the k-th smallest fraction that can be formed by dividing one element of a sorted array arr by another element later in the array. The array arr is strictly increasing, starts with 1, and all elements after the first are prime numbers. For every pair (i, j) where 0 <= i < j < arr.length, the fraction arr[i] / arr[j] is considered. We are to return the numerator and denominator of the k-th smallest fraction in the sorted list of all these fractions.

The input is a sorted list of unique integers with at least two elements and at most 1000 elements. The output is a pair of integers [arr[i], arr[j]] representing the fraction. Key constraints are that arr[0] == 1, elements after the first are prime, and all numbers are unique. We are guaranteed that 1 <= k <= arr.length * (arr.length - 1) / 2, so k is always valid.

Edge cases that could trip a naive implementation include the smallest array of length 2, where the only fraction is 1 / prime, and the largest array where a brute-force enumeration of all fractions could become too slow (up to ~500,000 fractions).

Approaches

The brute-force approach is straightforward. We generate all possible fractions arr[i]/arr[j] for i < j, store them in a list, sort them, and return the k-th smallest fraction. This works correctly because sorting all fractions guarantees the correct order. However, it is slow with time complexity O(n² log n²) and space complexity O(n²), which is inefficient for arrays of length 1000.

The optimal approach leverages a min-heap (priority queue) or binary search. We can think of the fractions in a matrix form where rows are numerators and columns are denominators. Using a min-heap, we only store the smallest possible fractions in each column initially and iteratively pop the smallest fraction while pushing the next fraction in the same column. This ensures we do not generate all fractions, reducing both time and space complexity to roughly O(k log n) and O(n), respectively.

Approach Time Complexity Space Complexity Notes
Brute Force O(n² log n²) O(n²) Generate all fractions and sort
Min-Heap Optimal O(k log n) O(n) Use heap to track next smallest fraction per column

Algorithm Walkthrough

  1. Initialize a min-heap. Each element is a tuple (fraction value, numerator index i, denominator index j).
  2. Push the first fraction from each row into the heap. That is, push (arr[i]/arr[-1], i, arr.length - 1) for i from 0 to arr.length - 2. This ensures we start with the smallest denominator in each row.
  3. Repeat k-1 times: pop the smallest fraction from the heap. Let the popped fraction have indices (i, j).
  4. If j-1 > i, push the next fraction from the same row (arr[i]/arr[j-1], i, j-1) into the heap. This moves left in the row to the next larger fraction.
  5. After k-1 pops, the fraction at the top of the heap is the k-th smallest fraction. Return [arr[i], arr[j]].

Why it works: Each row in the fraction matrix is sorted because numerators are fixed and denominators increase. The heap ensures the global smallest fraction is always on top. By pushing only the next fraction in the same row, we maintain the invariant that the heap always contains the next smallest unseen fractions.

Python Solution

from typing import List
import heapq

class Solution:
    def kthSmallestPrimeFraction(self, arr: List[int], k: int) -> List[int]:
        n = len(arr)
        min_heap = []
        
        # Initialize heap with fractions with the largest denominator
        for i in range(n - 1):
            heapq.heappush(min_heap, (arr[i] / arr[-1], i, n - 1))
        
        # Pop the smallest fraction k-1 times
        for _ in range(k - 1):
            _, i, j = heapq.heappop(min_heap)
            if j - 1 > i:
                heapq.heappush(min_heap, (arr[i] / arr[j - 1], i, j - 1))
        
        # The top of heap is the k-th smallest fraction
        _, i, j = heapq.heappop(min_heap)
        return [arr[i], arr[j]]

In this implementation, we initialize the heap with fractions that have the largest denominator. This ensures the heap contains the smallest fractions initially. Each pop gives the next smallest fraction, and we push the next fraction from the same numerator row if it exists. Finally, the k-th pop gives the answer.

Go Solution

package main

import (
    "container/heap"
)

type Fraction struct {
    value float64
    i     int
    j     int
}

type MinHeap []Fraction

func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i].value < h[j].value }
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *MinHeap) Push(x interface{}) {
    *h = append(*h, x.(Fraction))
}

func (h *MinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func kthSmallestPrimeFraction(arr []int, k int) []int {
    n := len(arr)
    h := &MinHeap{}
    heap.Init(h)
    
    for i := 0; i < n-1; i++ {
        heap.Push(h, Fraction{value: float64(arr[i]) / float64(arr[n-1]), i: i, j: n - 1})
    }
    
    for t := 0; t < k-1; t++ {
        f := heap.Pop(h).(Fraction)
        if f.j-1 > f.i {
            heap.Push(h, Fraction{value: float64(arr[f.i]) / float64(arr[f.j-1]), i: f.i, j: f.j - 1})
        }
    }
    
    f := heap.Pop(h).(Fraction)
    return []int{arr[f.i], arr[f.j]}
}

Go-specific considerations include using a struct for fractions and implementing the heap.Interface. Float64 is used for fraction comparison.

Worked Examples

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

Initial heap: [(1/5,0,3),(2/5,1,3),(3/5,2,3)]

1st pop: 1/5 (push 1/3): heap = [(1/3,0,2),(2/5,1,3),(3/5,2,3)]

2nd pop: 1/3 (push 1/2): heap = [(1/2,0,1),(2/5,1,3),(3/5,2,3)]

3rd pop: 2/5 (k-th fraction): return [2,5]

Example 2: arr = [1,7], k = 1

Initial heap: [(1/7,0,1)]

1st pop: 1/7: return [1,7]

Complexity Analysis

Measure Complexity Explanation
Time O(k log n) Each heap operation is O(log n), and we perform k pops/pushes
Space O(n) Heap stores at most n fractions at a time

The complexity is efficient because we avoid enumerating all O(n²) fractions.

Test Cases

# Basic examples
assert Solution().kthSmallestPrimeFraction([1,2,3,5], 3) == [2,5]  # Example 1
assert Solution().kthSmallestPrimeFraction([1,7], 1) == [1,7]     # Example 2

# Edge cases
assert Solution().kthSmallestPrimeFraction([1,2], 1) == [1,2]      # Minimum length
assert Solution().kthSmallestPrimeFraction([1,2,3,5,7,11,13], 10) == [5,11]  # Larger k
assert Solution().kthSmallestPrimeFraction([1,2,3,5,7,11,13], 1) == [1,13]   # Smallest fraction

# Stress test
arr = [1] + [i for i in range(2,1001) if all(i%p!=0 for p in range(2,int(i**0.5)+1))]
assert Solution().kthSmallestPrimeFraction(arr, len(arr)*(len(arr)-1)//2) == [arr[-2], arr[-1]]  # Largest k
Test Why
[1,2,3,5], k=3 Standard example to verify correct ordering
[1,7], k=1 Smallest array edge