중복된 값이 존재할 수 있는 정수 배열이 주어졌을 때, 특정 target 숫자에 해당하는 인덱스를 무작위로 하나 골라 반환하는 문제를 생각해 봅시다. 단, target은 반드시 배열 안에 존재한다고 가정합니다. 예를 들어 배열이 [1, 2, 3, 3, 3]이라면 pick(3)을 호출했을 때 인덱스 2, 3, 4 중 하나가 동등한 확률로 반환되어야 합니다.
접근 방법: 저수지 샘플링(Reservoir Sampling)
데이터의 크기를 미리 알지 못하거나 스트림 형태로 입력이 들어오는 상황에서도 균등한 확률로 샘플을 추출할 수 있는 대표적인 기법이 바로 저수지 샘플링입니다. 이 문제는 배열을 딱 한 번만 순회하면서(O(n)) 추가 메모리 없이(O(1)) 해결할 수 있다는 점이 큰 장점입니다.
알고리즘 단계
- ret := -1, cnt := 1로 초기화합니다.
- i를 0부터 배열의 크기까지 순회하며 다음을 반복합니다.
- v[i] == target인 경우, rand() % cnt == 0이면 ret = i로 갱신합니다.
- cnt를 1 증가시킵니다.
- 순회가 끝나면 ret을 반환합니다.
왜 동작할까?
target이 k번째로 등장할 때 해당 인덱스가 선택될 확률은 정확히 1/k입니다. 예를 들어 세 번째로 등장한 원소는 cnt = 3일 때 rand() % 3 == 0이 될 확률, 즉 1/3의 확률로 선택됩니다. 이후 새로운 후보가 계속 등장해도 기존 선택 결과가 덮어씌워질 확률을 곱해 계산해 보면, 모든 후보 인덱스가 최종적으로 동일한 1/n의 확률로 선택됨을 수학적으로 증명할 수 있습니다.
C++ 구현 예제
아래 구현을 통해 더 잘 이해해 보겠습니다.
#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
vector<int> v;
Solution(vector<int>& nums) {
srand(time(NULL)); // 난수 시드 초기화
v = nums;
}
int pick(int target) {
int ret = -1;
int cnt = 1;
for(int i = 0; i < v.size(); i++){
if(v[i] == target){
// k번째 등장 시 1/k 확률로 선택
if(rand() % cnt++ == 0) ret = i;
}
}
return ret;
}
};
main(){
vector<int> v = {1,2,3,3,3};
Solution ob(v);
cout << (ob.pick(3));
}
입력
[1,2,3,3,3]으로 초기화 pick(3)을 호출하여 무작위 인덱스 획득
출력
4 3 4 2
복잡도 분석
- 시간 복잡도: O(n) — 배열을 한 번만 순회합니다.
- 공간 복잡도: O(1) — 별도의 저장 공간이 필요하지 않습니다.
이처럼 저수지 샘플링을 활용하면 중복된 값이 여러 개 있더라도 어떤 인덱스든 공평하게 선택할 수 있으며, 데이터의 전체 크기를 미리 알지 못하는 상황에서도 매우 유용하게 적용할 수 있습니다.