Problem of The Day: Find K-th Smallest Pair Distance

Problem Statement

Brute Force - TLE

class Solution:
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        N = len(nums)
        min_heap = []
        for i in range(N - 1):
            for j in range(i + 1, N):
                distance = abs(nums[i] - nums[j])
                heapq.heappush(min_heap, distance)

        res = 0
        while k > 0 and min_heap:
            res = heapq.heappop(min_heap)
            k -= 1

        return res

class Solution:
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        nums.sort()
        N = len(nums)
        max_heap = []
        for i in range(N - 1):
            for j in range(i + 1, N):
                distance = abs(nums[i] - nums[j])
                heapq.heappush(max_heap, -distance)
                if len(max_heap) > k:
                    heapq.heappop(max_heap)

        return -max_heap[0]

Editorial

Approach 1: Bucket Sort - TLE

class Solution:
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        list_size = len(nums)

        # Find the maximum element in the list
        max_element = max(nums)

        # Initialize a bucket list to store counts of each distance
        distance_bucket = [0] * (max_element + 1)

        # Populate the bucket list with counts of each distance
        for i in range(list_size):
            for j in range(i + 1, list_size):
                # Compute the distance between nums[i] and nums[j]
                distance = abs(nums[i] - nums[j])

                # Increment the count for this distance in the bucket
                distance_bucket[distance] += 1

        # Find the k-th smallest distance
        for dist in range(max_element + 1):
            # Reduce k by the number of pairs with the current distance
            k -= distance_bucket[dist]

            # If k is less than or equal to 0, return the current distance
            if k <= 0:
                return dist

        return -1  # Return -1 if no distance found, should not reach here

Approach 2: Binary Search + Dynamic Programming (DP)

class Solution:
    def smallestDistancePair(self, nums, k):
        nums.sort()
        array_size = len(nums)

        # Highest element in the sorted array
        max_element = nums[-1]

        # Maximum possible distance
        max_possible_distance = max_element * 2

        # Initialize arrays for prefix counts and value counts
        prefix_count = [0] * max_possible_distance
        value_count = [0] * (max_element + 1)

        # Populate prefix count and value count
        index = 0
        for value in range(max_possible_distance):
            while index < array_size and nums[index] <= value:
                index += 1
            prefix_count[value] = index
        for i in range(array_size):
            value_count[nums[i]] += 1

        # Binary search to find kth smallest distance
        left, right = 0, max_element
        while left < right:
            mid = (left + right) // 2

            # Count pairs with distance <= mid
            count = self._count_pairs(nums, prefix_count, value_count, mid)

            # Adjust binary search bounds based on count
            if count < k:
                left = mid + 1
            else:
                right = mid
        return left

    def _count_pairs(self, nums, prefix_count, value_count, max_distance):
        count = 0
        array_size = len(nums)
        index = 0

        while index < array_size:
            current_value = nums[index]
            value_count_for_current = value_count[current_value]

            # Calculate pairs involving current value with distance <= max_distance
            pairs_with_larger_values = (
                prefix_count[
                    min(current_value + max_distance, len(prefix_count) - 1)
                ]
                - prefix_count[current_value]
            )
            pairs_with_same_values = (
                value_count_for_current * (value_count_for_current - 1) // 2
            )
            count += (
                pairs_with_larger_values * value_count_for_current
                + pairs_with_same_values
            )

            # Skip duplicate values
            while index + 1 < array_size and nums[index] == nums[index + 1]:
                index += 1
            index += 1

        return count

Approach 3: Binary Search + Sliding Window

class Solution:
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        nums.sort()
        array_size = len(nums)

        # Initialize binary search range
        low = 0
        high = nums[array_size - 1] - nums[0]

        while low < high:
            mid = (low + high) // 2

            # Count pairs with distance <= mid
            count = self._count_pairs_with_max_distance(nums, mid)

            # Adjust binary search bounds based on count
            if count < k:
                low = mid + 1
            else:
                high = mid

        return low

    # Count number of pairs with distance <= max_distance using a moving window
    def _count_pairs_with_max_distance(
        self, nums: List[int], max_distance: int
    ) -> int:
        count = 0
        array_size = len(nums)
        left = 0

        for right in range(array_size):
            # Adjust the left pointer to maintain the window with distances <=
            # max_distance
            while nums[right] - nums[left] > max_distance:
                left += 1
            # Add the number of valid pairs ending at the current right index
            count += right - left
        return count