Want Help Cracking FAANG?

(Then click this)

×
Back to Question Bank

786. K-th Smallest Prime Fraction - Leetcode Solution

Code Implementation

import heapq

class Solution:
    def kthSmallestPrimeFraction(self, arr, k):
        n = len(arr)
        heap = []
        # Min-heap: store (fraction value, numerator index, denominator index)
        for i in range(n-1):
            heapq.heappush(heap, (arr[i]/arr[n-1], i, n-1))
        for _ in range(k-1):
            val, i, j = heapq.heappop(heap)
            if j-1 > i:
                heapq.heappush(heap, (arr[i]/arr[j-1], i, j-1))
        _, i, j = heapq.heappop(heap)
        return [arr[i], arr[j]]
      
#include <vector>
#include <queue>
using namespace std;

class Solution {
public:
    vector<int> kthSmallestPrimeFraction(vector<int>& arr, int k) {
        int n = arr.size();
        auto cmp = [&arr](const vector<int>& a, const vector<int>& b) {
            return arr[a[0]] * arr[b[1]] > arr[b[0]] * arr[a[1]];
        };
        priority_queue<vector<int>, vector<vector<int>>, decltype(cmp)> pq(cmp);
        for (int i = 0; i < n - 1; ++i) {
            pq.push({i, n - 1});
        }
        for (int t = 0; t < k - 1; ++t) {
            auto frac = pq.top(); pq.pop();
            int i = frac[0], j = frac[1];
            if (j - 1 > i) {
                pq.push({i, j - 1});
            }
        }
        auto frac = pq.top();
        return {arr[frac[0]], arr[frac[1]]};
    }
};
      
import java.util.*;

class Solution {
    public int[] kthSmallestPrimeFraction(int[] arr, int k) {
        int n = arr.length;
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) ->
            arr[a[0]] * arr[b[1]] - arr[b[0]] * arr[a[1]]
        );
        for (int i = 0; i < n - 1; ++i) {
            pq.offer(new int[]{i, n - 1});
        }
        for (int t = 0; t < k - 1; ++t) {
            int[] frac = pq.poll();
            int i = frac[0], j = frac[1];
            if (j - 1 > i) {
                pq.offer(new int[]{i, j - 1});
            }
        }
        int[] frac = pq.poll();
        return new int[]{arr[frac[0]], arr[frac[1]]};
    }
}
      
class MinHeap {
    constructor(compare) {
        this.data = [];
        this.compare = compare;
    }
    push(val) {
        this.data.push(val);
        this._bubbleUp(this.data.length - 1);
    }
    pop() {
        if (this.data.length === 1) return this.data.pop();
        const top = this.data[0];
        this.data[0] = this.data.pop();
        this._bubbleDown(0);
        return top;
    }
    _bubbleUp(idx) {
        while (idx > 0) {
            let parent = Math.floor((idx - 1) / 2);
            if (this.compare(this.data[idx], this.data[parent]) < 0) {
                [this.data[idx], this.data[parent]] = [this.data[parent], this.data[idx]];
                idx = parent;
            } else break;
        }
    }
    _bubbleDown(idx) {
        let length = this.data.length;
        while (true) {
            let left = 2 * idx + 1, right = 2 * idx + 2, smallest = idx;
            if (left < length && this.compare(this.data[left], this.data[smallest]) < 0) smallest = left;
            if (right < length && this.compare(this.data[right], this.data[smallest]) < 0) smallest = right;
            if (smallest !== idx) {
                [this.data[idx], this.data[smallest]] = [this.data[smallest], this.data[idx]];
                idx = smallest;
            } else break;
        }
    }
}

var kthSmallestPrimeFraction = function(arr, k) {
    let n = arr.length;
    let heap = new MinHeap((a, b) => arr[a[0]] * arr[b[1]] - arr[b[0]] * arr[a[1]]);
    for (let i = 0; i < n - 1; ++i) {
        heap.push([i, n - 1]);
    }
    for (let t = 0; t < k - 1; ++t) {
        let [i, j] = heap.pop();
        if (j - 1 > i) heap.push([i, j - 1]);
    }
    let [i, j] = heap.pop();
    return [arr[i], arr[j]];
};
      

Problem Description

Given a sorted array arr containing prime numbers in increasing order, and an integer k, your task is to find the k-th smallest fraction of the form arr[i] / arr[j] where 0 ≤ i < j < arr.length.

Each fraction must use two distinct elements from arr (no repeated indices), and you are guaranteed there is exactly one valid answer for each input. Return the answer as an array [arr[i], arr[j]].

  • The array is sorted and contains only primes.
  • Each fraction is formed as arr[i] / arr[j] where i < j.
  • Return the k-th smallest fraction in ascending order of value.
  • There will always be one valid answer.

Thought Process

At first glance, the problem seems to require generating all possible fractions, sorting them, and picking the k-th smallest. This brute-force approach is straightforward but quickly becomes inefficient as the array grows, since the number of possible fractions is n * (n-1) / 2.

To optimize, we notice that since arr is sorted, the smallest fractions are those with the smallest numerators and largest denominators. We can use a min-heap (priority queue) to efficiently track the next smallest fraction without generating all possible fractions at once.

The key insight is to treat the set of possible fractions as a "matrix" where moving right or up gives larger fractions, and to always expand from the smallest unvisited fraction.

Solution Approach

  • Step 1: Initialize a min-heap (priority queue) to keep track of the next smallest fractions. Each entry in the heap contains the indices (i, j) representing the fraction arr[i] / arr[j].
  • Step 2: Push the smallest possible fractions into the heap. Since the array is sorted, for each i from 0 to n-2, we pair arr[i] with the largest denominator arr[n-1] and push (arr[i]/arr[n-1], i, n-1) into the heap.
  • Step 3: Pop the smallest fraction from the heap (this is the current smallest unvisited fraction). For each pop, if possible, push the next fraction in the same row (i.e., same numerator, next smaller denominator) into the heap.
  • Step 4: Repeat step 3 k-1 times. The next fraction popped from the heap is the k-th smallest fraction.
  • Step 5: Return the numerator and denominator corresponding to the k-th smallest fraction as an array.

This approach efficiently finds the k-th smallest fraction without sorting all possible fractions, thanks to the heap structure and the sorted nature of the array.

Example Walkthrough

Suppose arr = [1, 2, 3, 5] and k = 3.

  1. All possible fractions (with i < j): 1/2, 1/3, 1/5, 2/3, 2/5, 3/5.
  2. Sorted fractions: 1/5 (0.2), 1/3 (0.333...), 1/2 (0.5), 2/5 (0.4), 2/3 (0.666...), 3/5 (0.6).
  3. Order after sorting: 1/5, 1/3, 1/2, 2/5, 3/5, 2/3.
  4. The 3rd smallest is 1/2, so the answer is [1, 2].
  5. Using the heap approach:
    • Heap starts with: (1/5, 0, 3), (2/5, 1, 3), (3/5, 2, 3).
    • Pop 1/5, push 1/3 (0, 2).
    • Pop 1/3, push 1/2 (0, 1).
    • Pop 1/2 (third pop). Done!

The heap always gives the next smallest fraction, and after k pops, we have the answer.

Time and Space Complexity

  • Brute-Force Approach:
    • Generate all O(n^2) fractions.
    • Sort them: O(n^2 \log n) time.
    • Space: O(n^2) for storing all fractions.
  • Optimized Heap Approach:
    • Heap size is at most n at any time.
    • Each heap operation is O(\log n).
    • We do k pops and at most k pushes: O(k \log n) time.
    • Space: O(n) for the heap.

The heap approach is much faster and more space efficient, especially for large arr.

Summary

To solve the K-th Smallest Prime Fraction problem efficiently, we leverage the sorted nature of the input and use a min-heap to always expand the smallest unvisited fraction. This method avoids generating and sorting all possible fractions, resulting in a much faster and more elegant solution. The heap-based approach is a classic example of using data structures to optimize search problems involving sorted or partially ordered data.