[LeetCode] 719. Find K-th Smallest Pair Distance

https://leetcode.com/problems/find-k-th-smallest-pair-distance/description/?envType=daily-question&envId=2024-08-14

leetcode - Find K-th Smallest Pair Distance

문제 유형 : 투 포인터, 이진탐색

문제 난이도 : Hard

 

문제

The distance of a pair of integers a and b is defined as the absolute difference between a and b.

Given an integer array nums and an integer k, return the kth smallest distance among all the pairs nums[i] and nums[j] where 0 <= i < j < nums.length.

 

정수 a와 b의 거리는 두 수의 차의 절댓값이다.

정수 k와 정수 배열 nums가 주어진다.

k번째로 작은 nums[i]와 nums[j]의 거리를 구하시오.

 

풀이

처음에는 해시 맵에다가 [수 - 개수]의 꼴로 저장하여서 조합들을 만들어내고 그중 k번째를 찾아냈는데,

이는 중복되는 수가 거의 없으면 비효율적이다.

 

대신, 투 포인터와 이진탐색을 이용하여 효율적으로 해결할 수 있다.

우선, 배열을 정렬한다.

두 수의 거리의 범위를 [left, right]라 하자. left는 최소의 거리, 즉 0으로 시작하고, right는 최대의 거리, 즉 nums[n-1] - nums[0]으로 시작한다.

 

매번 중앙값 mid를 구하고, 투 포인터 i, j를 0으로 초기화한다.

0번째 인덱스부터 두 수의 차이가 mid이하인 동안 오른쪽 포인터 j를 넓혀서, mid이하의 거리의 쌍들의 개수를 누적시킨다.

 

만약 개수가 k보다 적다면, mid이하의 거리 쌍이 k개보다 작다는 뜻이므로, left를 mid+1로 하여 더 큰 거리를 탐색한다.

반대로, count가 k 이상이라면, 더 작은 거리 탐색을 위해 right을 mid로 한다.

 

코드

C++

 

1. 오답: 해시맵을 이용한 풀이(TC 18/19 Solved, 마지막 TC에서 TLE)

class Solution {
public:
    int smallestDistancePair(vector<int>& nums, int k) {
        int n = nums.size();
        sort(nums.begin(), nums.end());
        unordered_map<int, int> mp;

        for(const auto &num : nums) {
            mp[num]++;
        }

        unordered_map<int, int> mp2;
        
        for(const auto &element1 : mp) {
            for(const auto &element2 : mp) {
                int diff = abs(element1.first - element2.first);
                int amount = element1.first == element2.first? element1.second * (element1.second-1) : element1.second * element2.second;

                mp2[diff] += amount;
            }
        }

        vector<pair<int, int>> table;
        for(auto &element : mp2) {
            element.second /= 2;
            // cout << element.first << " " << element.second << endl;
            table.push_back(element);
        }

        sort(table.begin(), table.end());

        for(int i = 0; i < table.size(); i++) {
            k -= table[i].second;
            if(k <= 0) return table[i].first;
        }

        return -1;
    }
};

 

2. 정답: 이진 탐색과 투 포인터를 이용한 풀이

class Solution {
public:
    int smallestDistancePair(vector<int>& nums, int k) {
        sort(nums.begin(), nums.end());
        int n = nums.size();
        int left = 0;
        int right = nums[n-1] - nums[0];

        while(left < right) {
            int mid = (left + right) / 2;
            int count = 0, j = 0;

            for(int i = 0; i < n; i++) {
                while(j < n && nums[j] - nums[i] <= mid) j++;
                count += j - i - 1;
            }

            if(count < k) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }

        return left;
    }
};