주어진 집합에서 랜덤하게 원소를 고르기

어떤 유한한 집합이 주어지고, 여기에서 랜덤하게 중복없이 n개의 원소를 골라내는 작업을 구현해보자. 가장 간단하게 생각한다면 다음과 같은 알고리듬을 생각할 수 있다.

 

1. 먼저 골라낼 원소를 가리키는 난수값을 발생시킨다.
2. 1의 난수가 가리키는 집합의 원소를 고른다.
3. 이 때 고르게 된 값이 이전에 뽑은 값과 중복되는지를 체크한다. 만약 한 번이라도 뽑았던 값이라면 버리고, 그렇지 않은 경우 결과에 추가한다.
4. 1~3의 과정을 결과 집합의 크기가 n이 될때까지 반복한다.

로또 번호는 1~45 사이의 범위가 될 수 있기 때문에 random.randint(1, 45)를 통해서 랜덤하게 고를 수 있다.

from random import randint

result = []
while len(result) < 6:
  a = randint(1, 45)
  if a not in result:
    result.append(a)

이 때 result 가 리스트라면 in 연산은 O(n)의 복잡도를 가진다. 만약 result가 set 타입이라면 in 연산이 O(1)이 될 것이므로 조금의 성능 향상을 기대할 수 있다.

 

그런데 굳이 이렇게 반복을 직접 하지 않더라도 random 모듈에는 sample 이라는 함수가 존재한다. sample 함수는 주어진 집합에서 n 개의 원소를 중복없이 샘플링한다.(단, 같은 원소를 두 번 뽑지 않을 뿐이지 원래 집합에 중복된 원소가 있다면 결과에도 중복된 값이 들어갈 수 있다.)

from random import sample
xs = sample(range(1, 46), 6)
# xs = [15, 2, 29, 42, 9, 22]

조금 더 깊이

이번에는 뽑게되는 숫자에 특정한 규칙을 추가해보자. 뽑히는 숫자 6개 중에서 홀수/짝수의 비율이 3:3 이거나 2:4 혹은 4:2 가 되도록 하는 것이다. 이를 위해서는 다음과 같이 샘플링을 하면 된다.

 

1. 먼저 홀수, 짝수는 각각 최소 2개씩은 샘플링된다.
2. 1에서 4개를 뽑아 놓은 뒤, 홀/짝 구분 없이 전체 집합에서 중복 없이 2개를 샘플링한다.
3. 1, 2의 합집합이 최종 결과가 된다.

따라서 홀수, 짝수의 집합을 set으로 만들고, 여기서 샘플링한 결과를 다시 set으로 만들어서 채워나간다.

from random import sample

evens = set(range(2, 46, 2))
odds = set(range(1, 46, 2))
x = set(sample(evens, 2))
x |= set(sample(odds, 2))
x |= set(sample((odds | evens) - x, 2))

set을 사용하면 |(합집합), &(교집합), -(차집합) 및 ^(대칭차집합) 연산을 사용할 수 있다. 이 때 각 연산은 set 객체의 .union(), .intersection(), .difference(), .symmetric_difference()에 해당한다.