719. Find K-th Smallest Pair Distance

[zeikar]

Problem link

https://leetcode.com/problems/find-k-th-smallest-pair-distance/

Problem Summary

배열 중에서 pair의 차이가 k번째로 작은 값을 찾는 문제

Solution

대충 감으로 이진 탐색으로 찾아야할 것 같다. 그런데 어떻게 n보다 작은 pair 개수를 셀 수 있을까?
슬라이딩 윈도우로 셀 수 있다.

먼저 배열을 정렬해주자.
그 후 n보다 작은 pair의 개수를 세는 함수를 정의하고 n을 이진 탐색으로 돌면서 k개가 나올 때까지 탐색해주면 된다.

n보다 작은 pair의 개수는 슬라이딩 윈도우로 가능한데 윈도우의 양쪽 차이가 n이하라면 윈도우의 크기가 pair의 개수가 된다. 왜냐하면 윈도우 내부 pair들은 전부 n보다 작기 때문이다.
죽 돌면서 개수를 세고 리턴 후 이진 탐색을 계속해주면 된다.

에디토리얼 참고

Source Code

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

        def countSmallerPair(n):
            ret = 0
            left = 0
            for right in range(len(nums)):
                while left <= right and nums[right] - nums[left] > n:
                    left += 1

                ret += right - left
            return ret

        s, e = 0, nums[-1] - nums[0]
        while s < e:
            mid = s + (e - s) // 2

            cnt = countSmallerPair(mid)
            if cnt < k:
                s = mid + 1
            else:
                e = mid
        return s