C++로 배열에서 고유한 K-diff 쌍 개수 구하는 방법

배열과 정수 k가 주어졌을 때, 배열 안에 존재하는 고유한 k-diff 쌍의 개수를 구하는 문제입니다. 여기서 k-diff 쌍이란 (i, j) 형태를 가지며, i와 j가 모두 배열에 존재하고 두 값의 절댓값 차이가 k와 같은 경우를 의미합니다.

예를 들어 입력이 [3, 1, 4, 1, 5]이고 k = 2라고 가정해 보겠습니다. 이 경우 출력은 2가 됩니다. 배열에는 (1, 3)과 (3, 5)라는 두 개의 2-diff 쌍이 존재하기 때문입니다.

문제 해결 접근 방법

이 문제는 해시 맵(map)과 집합(set)을 활용하면 효율적으로 해결할 수 있습니다. 단계별로 살펴보겠습니다.

seen과 done이라는 두 개의 맵을 선언합니다.
집합(set) s를 하나 선언합니다.
k < 0이면 절댓값 차이가 음수일 수 없으므로 0을 반환합니다.
배열을 순회하면서 각 숫자의 등장 횟수를 seen에 기록하고, 해당 값을 집합 s에 삽입하여 중복을 제거합니다.
정답 변수 ans를 0으로 초기화합니다.
집합 s의 각 원소 it에 대해 다음을 수행합니다.
- k == 0인 경우: 동일한 값으로 이루어진 쌍을 찾아야 하므로, seen[it]이 1보다 크다면(중복된 값이 존재한다면) ans를 1 증가시킵니다.
- 그 외의 경우: 먼저 done[it]을 1 증가시켜 현재 처리 중임을 표시합니다. 그런 다음 (it + k)가 seen에는 존재하지만 done에는 없다면 ans를 증가시키고, 마찬가지로 (it - k)가 seen에는 존재하지만 done에는 없다면 ans를 증가시킵니다. 이 과정을 통해 동일한 쌍이 중복해서 계산되는 것을 방지할 수 있습니다.
모든 순회가 끝나면 ans를 반환합니다.

C++ 구현 예제

아래 코드를 통해 더 자세히 이해해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    int findPairs(vector<int>& nums, int k) {
        map<int, int> seen, done;
        set<int> s;
        if (k < 0)
            return 0;
        for (int i = 0; i < nums.size(); i++) {
            seen[nums[i]]++;
            s.insert(nums[i]);
        }
        int ans = 0;
        for (auto it = s.begin(); it != s.end(); it++) {
            if (k == 0) {
                if (seen[*it] > 1)
                    ans++;
            }
            else {
                done[*it]++;
                if (seen.find(*it + k) != seen.end() && done.find(*it + k) == done.end())
                    ans++;
                if (seen.find(*it - k) != seen.end() && done.find(*it - k) == done.end())
                    ans++;
            }
        }
        return ans;
    }
};
main(){
    Solution ob;
    vector<int> v = {3,1,4,1,5};
    cout << (ob.findPairs(v, 2));
}

입력

{3,1,4,1,5}, 2

출력

복잡도 분석

이 알고리즘은 배열을 한 번 순회하여 맵과 집합을 채우고, 고유한 값의 개수만큼 다시 순회하므로 시간 복잡도는 O(n log n)입니다. 맵과 집합이 내부적으로 균형 이진 탐색 트리를 사용하기 때문입니다. 공간 복잡도 역시 저장되는 고유한 값의 개수에 비례하여 O(n)입니다.