양의 정수 배열 w가 주어졌을 때, 각 원소 w[i]는 인덱스 i의 가중치를 의미합니다. 이때 가중치에 비례하여 인덱스를 무작위로 선택하는 함수 pickIndex()를 구현해야 합니다.
예를 들어 입력이 [1, 3]이라면, 인덱스 0은 전체 가중치의 1/4 확률로, 인덱스 1은 3/4 확률로 선택됩니다. 따라서 pickIndex()를 다섯 번 호출하면 결과가 다음과 같이 나올 수 있습니다.
0, 1, 1, 1, 0
해결 접근 방법: 누적 합(Prefix Sum) 활용
이 문제는 누적 합(prefix sum)과 이진 탐색(binary search)을 조합하면 효율적으로 해결할 수 있습니다. 단계별로 살펴보겠습니다.
- 생성자에서 배열 w의 각 원소를 이전 원소와 누적하여 더해 누적 합 배열 v를 만듭니다.
- 예를 들어
[1, 3]은 누적 합으로 변환하면[1, 4]가 됩니다. pickIndex()가 호출되면rand() % v.back()으로 0부터 마지막 누적 값 미만 범위의 난수 r을 생성합니다.- 누적 합 배열 v에서 r보다 큰 첫 번째 값을
upper_bound()로 찾아 그 위치(인덱스)를 반환합니다.
이 방식의 핵심 아이디어는 누적 합 배열을 수직선 위의 구간으로 생각하는 것입니다. 가중치가 클수록 해당 인덱스가 차지하는 구간이 넓어지므로, 난수가 어느 구간에 속하는지 확인하면 자연스럽게 가중치에 비례한 선택이 이루어집니다.
C++ 구현 예제
#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
int n;
vector <int> v;
Solution(vector<int>& w) {
srand(time(NULL));
n = w[0];
for(int i = 1; i < w.size(); i++){
w[i] += w[i - 1];
n = w[i];
}
v = w;
}
int pickIndex() {
return upper_bound(v.begin(), v.end(), rand() % v.back()) - v.begin();
}
};
main(){
vector<int> v = {1,3};
Solution ob(v);
cout << (ob.pickIndex()) << endl;
cout << (ob.pickIndex()) << endl;
cout << (ob.pickIndex()) << endl;
cout << (ob.pickIndex()) << endl;
cout << (ob.pickIndex()) << endl;
}입력
[1, 3]으로 초기화한 뒤 pickIndex()를 다섯 번 호출합니다.
출력
1 1 1 1 0
복잡도 분석
- 시간 복잡도: 생성자에서 누적 합 계산에 O(n),
pickIndex()는 이진 탐색을 사용하므로 호출당 O(log n)입니다. - 공간 복잡도: 누적 합 배열을 저장하기 위해 O(n)의 추가 공간이 필요합니다.
매번 선형 탐색으로 구간을 찾는 O(n) 방식과 비교했을 때, 이진 탐색을 활용하면 호출 횟수가 많아지는 상황에서도 성능을 크게 개선할 수 있다는 점이 이 구현의 장점입니다.