LeetCode 719: Find K-th Smallest Pair Distance

All pairs with max heap

Generate all possible pairs keeping $k$ smallest in a max heap.

Time: $\mathcal{O}(n^2 \cdot \lg{k})$ , space: $\mathcal{O}(k)$

from heapq import heappush as push, heappop as pop

class Solution:
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        maxq = []
        n = len(nums)
        for i in range(n-1):
            for j in range(i+1, n):
                dist = abs( nums[i] - nums[j] )
                push( maxq, -dist )
                if len(maxq) > k:
                    pop(maxq)
        return -maxq[0]

Sorting and minq

We want to generate the pairs in sorted order of distance. Sort nums, put $(n-1)$ pairs having different first elements in a minq. These are like heads of sorted linked lists. Now pop $(k-1)$ times and return the top distance.

Time: $\mathcal{O}(\max{(k, n)} \cdot \lg{n})$ , space: $\mathcal{O}(n)$ .

from heapq import heappush as push, heappop as pop

class Solution:
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        n = len(nums)
        nums.sort()
        minq = []
        for i in range(n-1):
            dist = abs( nums[i]-nums[i+1] )
            push( minq, ( dist, (i, i+1) ) )
        
        while k > 1:
            k -= 1

            d, coords = pop( minq )
            i, j = coords
            if j == n-1:
                continue
            j += 1
            d = abs( nums[i]-nums[j] )
            push( minq, (d, (i, j)) )

        kth, _ = minq[0]

        return kth

Reformulate and binary search

If we have a monotonic function like $f(x)$ , we can find the minimum $x$ such that $f(x) \ge k$ using binary-search. We use mid as the $x$ and depending on whether $f(x) \ge k$ , we can move left or right.

Let $P\left(d_m\right)$ be a function that counts the pairs in nums having distance at most $d_m$ . This is a monotonic function because, if we increase $d_m$ , $P\left(d_m\right)$ either increases or remains the same, but never decreases. In other words, as we increase $d_m$ , if we keep asking, “Is $P\left( d_m \right) \ge k$ “, we have the sequence of answers like: $\langle \texttt{false}, \texttt{false}, \ldots, \texttt{false}, \texttt{true}, \texttt{true}, \ldots \rangle$ . So, the boolean version $\text{Boolean-P}\left(d_m, k\right)$ behaves like a step function.

The time complexity of the binary-search is $\mathcal{O}( \texttt{time}(\text{Boolean-P}) \cdot \lg{ \left|\text{Domain}(d_m) \right| } )$ .

For this problem, if we have sorted nums, we can find the count of pairs with a given maximum distance of $d_m$ in $\mathcal{O}(n)$ time.

Time complexity: $\mathcal{O}( n \cdot \lg{ \left( \max{(n, D)} \right) } )$ where $D$ is the distance between smallest and largest elements in nums. Space: that of sorting.

class Solution:
    def has_kth_smallest(self, nums, k, max_distance) -> bool:
        n = len(nums)
        start, pair_count = 0, 0
        for end in range(n):
            # If no pairs with this end,
            # start == end, we get pair_count += 0
            while abs( nums[end]-nums[start] ) > max_distance:
                start += 1
            pair_count += (end-start)
 
        return pair_count >= k
 
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        n = len(nums)
        # For efficient has_kth_smallest()
        nums.sort()
         
        lo, hi = 0, nums[-1]-nums[0]
        kth = -1
        while lo <= hi:
            mid = (lo+hi) // 2
            if self.has_kth_smallest(nums, k, mid):
                kth = mid
                hi = mid-1
            else:
                lo = mid+1
        return kth

unreasonably effective