문제 개요
이번 글에서는 모든 연산이 O(1) 시간 안에 수행되어야 하는 자료구조를 설계해 보겠습니다. 이 자료구조는 다음 세 가지 연산을 지원해야 합니다.
- insert(x): 컬렉션에 x를 삽입합니다.
- remove(x): 컬렉션에서 x를 삭제합니다.
- getRandom(): 컬렉션에 있는 원소 중 하나를 무작위로 반환합니다.
중복 허용 버전이므로 같은 값이 여러 번 삽입될 수 있으며, 각 값은 동일한 확률로 랜덤 선택되어야 합니다.
핵심 아이디어: 동적 배열 + 해시 맵 조합
이 문제를 효율적으로 풀기 위해서는 두 가지 자료구조를 함께 사용합니다.
- 동적 배열(nums): 실제 데이터를 {값, 해당 값의 인덱스 위치} 쌍으로 저장합니다. 배열은 인덱스 접근이 O(1)이고, 마지막 원소의 추가·삭제도 O(1)입니다.
- 해시 맵(m): 각 값이 배열 내 어느 위치에 있는지 그 인덱스 목록을 저장합니다.
삭제 시에는 대상 원소와 배열의 마지막 원소를 서로 교환(swap)한 뒤 마지막 원소만 제거합니다. 이렇게 하면 배열 중간의 원소를 삭제하더라도 O(1) 시간을 유지할 수 있습니다.
알고리즘 단계
- {값, 인덱스} 쌍을 저장할 배열 nums를 준비합니다.
- 각 값의 위치 정보를 담을 해시 맵 m을 준비합니다.
- insert(val) 함수를 정의합니다.
- ret := val이 m에 없으면 true (새로운 값인지 여부)
- m[val]의 끝에 nums의 현재 크기를 추가합니다.
- nums의 끝에 {val, m[val]의 크기 - 1} 쌍을 추가합니다.
- ret을 반환합니다.
- remove(val) 함수를 정의합니다.
- ret := val이 m에 있으면 true
- ret이 참이라면:
- last := nums의 마지막 원소
- m[last.first][last.second] := m[val]의 마지막 값 (마지막 원소의 위치 정보를 삭제 대상 위치로 갱신)
- nums[m[val].back()] := last (배열에서도 교환 반영)
- m[val]에서 마지막 원소를 제거합니다.
- m[val]이 비었다면 m에서 val을 완전히 삭제합니다.
- nums에서 마지막 원소를 제거합니다.
- ret을 반환합니다.
- getRandom() 함수를 정의합니다.
- nums에서 무작위 인덱스 하나를 골라 해당 값을 반환합니다.
배열의 마지막 원소를 삭제 대상 위치로 옮긴 후 제거하는 방식 덕분에, 중복된 값이 여러 개 있어도 각 연산의 시간 복잡도가 O(1)로 유지됩니다.
C++ 구현 예제
#include <bits/stdc++.h>
using namespace std;
class RandomizedCollection {
public:
vector <pair <int, int>> nums;
unordered_map <int, vector<int>> m;
RandomizedCollection() {
}
bool insert(int val) {
bool ret = m.find(val) == m.end();
m[val].push_back(nums.size());
nums.push_back({val, m[val].size() - 1});
return ret;
}
bool remove(int val) {
bool ret = m.find(val) != m.end();
if(ret){
pair <int, int> last = nums.back();
m[last.first][last.second] = m[val].back();
nums[m[val].back()] = last;
m[val].pop_back();
if(m[val].empty())m.erase(val);
nums.pop_back();
}
return ret;
}
int getRandom() {
return nums[rand() % nums.size()].first;
}
};
main(){
RandomizedCollection ob;
ob.insert(10);
ob.insert(35);
ob.insert(20);
ob.insert(40);
cout << (ob.getRandom()) << endl;
ob.remove(20);
cout << (ob.getRandom()) << endl;
}입력 예시
10, 35, 20, 40을 차례로 삽입한 뒤 랜덤 원소를 하나 얻습니다(예: 40). 이후 20을 삭제하고 다시 랜덤 원소를 얻습니다(예: 35).
출력 결과
40 35
복잡도 분석
- 시간 복잡도: insert, remove, getRandom 모두 평균 O(1)입니다. 해시 맵의 연산과 배열의 마지막 원소 접근·삭제가 상수 시간에 이루어지기 때문입니다.
- 공간 복잡도: O(N). 저장된 원소 수에 비례하여 배열과 해시 맵이 공간을 사용합니다.