이번 글에서는 평균 O(1) 시간 복잡도로 다음 세 가지 연산을 모두 지원하는 자료구조를 파이썬으로 구현해 보겠습니다.
- insert(val) – 집합에 값 val이 존재하지 않으면 삽입합니다.
- remove(val) – 집합에 값 val이 존재하면 제거합니다.
- getRandom() – 현재 집합에 있는 요소 중 하나를 무작위로 반환합니다. 이때 모든 요소가 동일한 확률로 선택되어야 합니다.
문제 접근 방법
단순히 배열만 사용하면 삭제 시 요소를 찾는 데 O(n)이 걸리고, 딕셔너리만 사용하면 랜덤 추출이 어렵습니다. 따라서 두 자료구조를 조합하는 것이 핵심입니다. 해결 과정은 다음과 같습니다.
- 초기화 단계에서 값의 존재 여부를 저장할 딕셔너리(present)와 실제 요소를 담을 배열(elements)을 준비합니다.
- insert(val): val이 딕셔너리에 없거나 present[val]이 0이라면, 배열 끝에 val을 추가하고 present[val]을 1로 설정한 뒤 True를 반환합니다. 이미 존재하는 값이면 False를 반환합니다.
- remove(val): val이 없거나 present[val]이 0이면 False를 반환합니다. 그렇지 않으면 present[val]을 0으로 만들고, 배열에서 val의 인덱스를 찾습니다.
- 찾은 인덱스가 마지막 위치가 아니라면, 마지막 요소를 해당 위치로 옮겨와 교체(swap)합니다. 이렇게 하면 배열 끝에서 요소를 제거할 수 있어 O(1) 삭제가 가능합니다.
- 마지막으로 배열의 마지막 요소를 pop()으로 제거하고 True를 반환합니다.
- getRandom(): 배열에 남아 있는 요소 중 하나를 random.choice()로 반환합니다. 배열에는 유효한 요소만 남아 있으므로 균등한 확률이 보장됩니다.
파이썬 구현 예제
아래 코드를 통해 전체 동작을 확인해 보겠습니다.
import random
class RandomizedSet(object):
def __init__(self):
self.present = {}
self.elements = []
def insert(self, val):
if val not in self.present or self.present[val] == 0:
self.elements.append(val)
self.present[val] = 1
return True
return False
def remove(self, val):
if val not in self.present or self.present[val] == 0:
return False
self.present[val] = 0
index = self.elements.index(val)
if index != len(self.elements)-1:
temp = self.elements[-1]
self.elements[-1] = val
self.elements[index] = temp
self.elements.pop()
return True
def getRandom(self):
return random.choice(self.elements)
ob = RandomizedSet()
print(ob.insert(1))
print(ob.remove(2))
print(ob.insert(2))
print(ob.getRandom())
print(ob.remove(1))
print(ob.insert(2))
print(ob.getRandom())
입력
클래스를 초기화한 후 insert(), remove(), getRandom() 함수를 순서대로 호출합니다.
출력
True
False
True
2
True
False
2
동작 설명 및 시간 복잡도
실행 결과를 단계별로 살펴보면 다음과 같습니다.
- insert(1) → 1이 없으므로 삽입, True
- remove(2) → 2가 존재하지 않으므로 False
- insert(2) → 2를 삽입, True
- getRandom() → {1, 2} 중 하나 반환, 2
- remove(1) → 1을 제거, True
- insert(2) → 2가 이미 있으므로 False
- getRandom() → {2}만 남았으므로 항상 2
삽입과 랜덤 추출은 명확하게 O(1)입니다. 삭제는 딕셔너리 조회가 O(1), 배열 내 위치 탐색(index)이 최악의 경우 O(n)이지만, 값을 마지막 요소와 교체한 뒤 pop()으로 제거하는 방식 덕분에 평균적으로 O(1)에 가깝게 동작합니다. 참고로 딕셔너리에 값의 인덱스를 함께 저장하면 index() 호출을 없애 엄밀한 O(1) 삭제도 구현할 수 있습니다.