문제 개요
양의 정수로 이루어진 배열 A가 있다고 가정해 보겠습니다. 어떤 부분 배열(연속된 구간) 안에 포함된 서로 다른 정수의 개수가 정확히 K라면, 그 부분 배열을 "좋은(good) 부분 배열"이라고 부릅니다. 예를 들어 배열 [1,2,3,1,2]에는 1, 2, 3이라는 세 가지 서로 다른 정수가 존재합니다. 우리의 목표는 배열 A에서 좋은 부분 배열이 총 몇 개인지 구하는 것입니다.
예를 들어 입력이 [1,2,3,1,4]이고 K = 3일 때 출력은 4가 됩니다. 정확히 세 개의 서로 다른 정수를 포함하는 부분 배열은 다음 네 가지이기 때문입니다.
[1,2,3], [1,2,3,1], [2,3,1], [3,1,4]
접근 방식: 슬라이딩 윈도우
이 문제를 효율적으로 해결하기 위한 핵심은 "서로 다른 정수가 최대 k개인 부분 배열의 개수"를 구하는 atMost() 함수를 먼저 정의하는 것입니다. 그다음, 정확히 K개의 서로 다른 정수를 가지는 부분 배열의 개수는 아래 식으로 계산할 수 있습니다.
atMost(a, k) - atMost(a, k - 1)
atMost() 함수는 다음과 같은 단계로 동작합니다.
- 윈도우의 오른쪽 끝을 가리키는 포인터 i, 왼쪽 끝을 가리키는 포인터 j를 사용하고, 각 숫자의 등장 횟수를 저장할 해시 맵 m을 준비합니다.
- i를 한 칸씩 이동하면서 a[i]를 윈도우에 추가합니다. 이때 처음 등장한 새로운 숫자라면 남은 허용치 k를 1씩 줄입니다.
- k가 음수가 되면, 즉 윈도우 내 고유 숫자가 k개를 초과하면 j를 이동시켜 왼쪽부터 원소를 제거합니다. 어떤 숫자가 윈도우에서 완전히 사라질 때마다 k를 1씩 되돌립니다.
- 각 i 위치에서 유효한 부분 배열의 개수는 현재 윈도우 길이인 (i - j + 1)이므로, 이 값을 매번 정답 ans에 더해 줍니다.
- 배열 전체를 순회한 뒤 ans를 반환합니다.
C++ 구현 예제
아래 코드를 통해 실제 구현을 더 잘 이해할 수 있습니다.
#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
int subarraysWithKDistinct(vector<int>& a, int k) {
return atMost(a, k) - atMost(a, k - 1);
}
int atMost(vector <int>& a, int k){
set <int> current;
int j = 0;
int ans = 0;
int n = a.size();
unordered_map <int, int> m;
for(int i = 0; i < a.size(); i++){
if(!m[a[i]]++) k--;
while(k < 0){
if(!--m[a[j]])
k++;
j++;
}
int x = ((i - j) + 1);
ans += x;
}
return ans;
}
};
main(){
Solution ob;
vector<int> v = {1,2,3,1,4};
cout << (ob.subarraysWithKDistinct(v, 3));
}입력
{1,2,3,1,4}, 3출력
4
시간 및 공간 복잡도
두 포인터 i와 j가 각각 배열을 한 번씩만 순회하므로 시간 복잡도는 O(n)입니다. 또한 숫자별 등장 횟수를 저장하는 해시 맵에 최대 n개의 키가 저장될 수 있으므로 공간 복잡도 역시 O(n)입니다. 덕분에 브루트포스 방식(O(n²))보다 훨씬 효율적으로 큰 입력까지 처리할 수 있습니다.