· 개발

길이를 모르는 스트림에서 k개를 고르게 뽑기: 저수지 샘플링(Algorithm R)

전체 개수를 모르는 데이터를 한 번만 읽으면서 k개를 무작위로 뽑는 저수지 샘플링. 교체 확률이 왜 k/i인지, 모든 원소가 k/n으로 남는 증명, k개 조합 전체가 고른지, 흔한 범위 실수까지 정리했습니다.

개발알고리즘

몇 GB짜리 로그에서 무작위로 100줄만 뽑아 보고 싶다고 해 봅시다. 줄 수는 모릅니다. 먼저 끝까지 세고 다시 읽으며 뽑으면 파일을 두 번 읽어야 하고, 표준 입력이나 네트워크 스트림처럼 다시 읽을 수 없는 입력이면 그마저 안 됩니다. 전부 메모리에 올린 뒤 섞는 방법은 크기가 커질수록 부담입니다.

저수지 샘플링(reservoir sampling)은 이 상황을 위한 방법입니다. 입력을 한 번만 읽고, 메모리는 k칸짜리 배열 하나만 씁니다. 입력이 언제 끝나든 그 순간 배열에 든 k개가 무작위 표본입니다.

Algorithm R

NIST의 Dictionary of Algorithms and Data Structures(DADS)는 이 알고리즘을 다음처럼 정의합니다.

  1. 처음 k개는 그대로 크기 k인 배열(저수지)에 넣는다.
  2. j번째 원소(j > k)가 오면 1부터 j까지의 정수 M을 무작위로 고른다.
  3. M ≤ k이면 배열의 M번째 칸을 j번째 원소로 바꾼다. 아니면 버린다.

j번째 원소가 저수지에 들어갈 확률은 k/j이고, 들어간다면 어느 칸을 밀어낼지는 k칸 중에서 고르게 정해집니다. M 하나로 두 가지를 한 번에 정하는 셈입니다. DADS는 이 방법이 Knuth의 Algorithm 3.4.2R을 메모리 안에서 처리하도록 바꾼 형태이고, Knuth는 이를 Alan G. Waterman의 것으로 밝혔다고 적고 있습니다. 이 글에서도 줄여서 Algorithm R이라고 부릅니다.

JavaScript로 옮기면 다음과 같습니다. 인덱스는 0부터 시작하므로 M 대신 0 이상 seen 미만의 m을 뽑고, m < k이면 바꿉니다.

// 0 이상 n 미만의 정수를 고르게 뽑는다 (1 ≤ n ≤ 2^32)
function randomBelow(n) {
  const limit = Math.floor(2 ** 32 / n) * n;
  const buf = new Uint32Array(1);
  do {
    crypto.getRandomValues(buf);
  } while (buf[0] >= limit);
  return buf[0] % n;
}

function reservoirSample(iterable, k) {
  if (!Number.isInteger(k) || k < 0) {
    throw new RangeError("k는 0 이상의 정수여야 합니다");
  }
  const reservoir = [];
  if (k === 0) return reservoir;

  let seen = 0;
  for (const item of iterable) {
    seen++;
    if (seen <= k) {
      reservoir.push(item);
      continue;
    }
    if (seen > 2 ** 32) throw new RangeError("randomBelow의 범위를 넘었습니다");
    const m = randomBelow(seen); // 0 ≤ m < seen
    if (m < k) reservoir[m] = item;
  }
  return reservoir;
}

randomBelow는 32비트 난수 가운데 floor(2^32 / n) × n 이상인 값을 버리고 다시 뽑습니다. 남는 값의 개수가 n의 배수가 되므로 % n 결과가 0부터 n−1까지 같은 횟수로 대응됩니다. 이 함수는 n이 2^32를 넘으면 쓸 수 없으니, 원소가 그보다 많을 수 있는 입력이라면 더 넓은 범위의 정수 난수가 필요합니다.

시간은 원소 하나당 상수 작업이라 O(n)이고, 추가 메모리는 O(k)입니다. 난수는 k번째 이후 원소마다 하나씩 뽑으므로, k ≥ 1이면 모두 max(n − k, 0)번입니다. 입력이 k개 이하면 한 번도 뽑지 않고, k = 0이면 입력을 읽기 전에 돌아가므로 역시 0번입니다.

왜 모든 원소가 k/n 확률로 남는가

입력이 n개에서 끝났고 1 ≤ k ≤ n이라고 합시다. 다음은 이 글의 계산입니다. (k = 0이면 함수가 빈 배열을 돌려주므로, n > 0일 때 모든 원소가 남을 확률은 0입니다. n = 0이면 확률을 따질 원소가 없습니다.)

먼저 j번째 원소(j > k)를 처리하는 단계에서 저수지의 특정 칸이 밀려날 확률을 봅니다. M이 정확히 그 칸 번호여야 하므로 1/j입니다. 그 칸에 있던 원소가 이 단계를 버틸 확률은 1 − 1/j = (j − 1)/j입니다.

처음 k개 중 하나는 확률 1로 들어간 뒤 k+1번째부터 n번째 단계까지 모두 버텨야 합니다.

(k/(k+1)) × ((k+1)/(k+2)) × … × ((n−1)/n) = k/n

분자와 분모가 차례로 약분되어 k/n만 남습니다.

**i번째 원소(i > k)**는 먼저 확률 k/i로 들어가고, i+1번째부터 n번째 단계까지 버텨야 합니다. 버틸 확률은 같은 방식으로 i/n이 됩니다.

(k/i) × (i/n) = k/n

어느 원소든 최종 저수지에 남을 확률은 k/n으로 같습니다. 입력이 어디서 끝나든 성립하므로 n을 미리 알 필요가 없습니다.

원소별 확률이 같다고 끝이 아니다

원소마다 k/n이라는 사실만으로는 “고르게 뽑았다”고 하기에 부족합니다. 예를 들어 원소 4개를 {1, 2}와 {3, 4} 두 묶음으로 나누고 동전을 던져 한 묶음을 고르면, 각 원소가 뽑힐 확률은 모두 1/2입니다. 하지만 {1, 3} 같은 조합은 절대 나오지 않습니다.

DADS는 Algorithm R에서 가능한 k개 조합 각각이 같은 확률로 나온다고 적고 있습니다. 이것도 귀납법으로 확인할 수 있습니다. 이번에도 이 글의 계산이고, 아래에서 k로 나누는 단계가 있으므로 1 ≤ k ≤ n을 가정합니다. k = 0이면 가능한 조합은 빈 집합 하나뿐이고 위 코드도 항상 빈 배열을 돌려주므로 따로 증명할 것이 없습니다. t개(t ≥ k)를 처리한 뒤 처음 t개에서 고를 수 있는 k개 조합 C(t, k)가지가 모두 1/C(t, k) 확률이라고 가정하고, t+1번째 원소를 처리합니다. 아래 계산에는 C(t+1, k) = C(t, k) × (t+1)/(t+1−k)를 씁니다.

  • 새 원소가 없는 조합 S는 직전 저수지가 S였고 새 원소가 버려졌을 때만 나옵니다. 확률은 (1/C(t, k)) × (t+1−k)/(t+1) = 1/C(t+1, k)입니다.
  • 새 원소가 있는 조합 S에서 새 원소를 뺀 k−1개를 T라고 합시다. 직전 저수지는 T에 원소 x 하나를 더한 것이어야 하고, x의 후보는 t − (k−1)개입니다. 새 원소가 들어오고(k/(t+1)) x의 칸이 골라져야(1/k) 하므로 전체 확률은 (t−k+1) × (1/C(t, k)) × (1/(t+1)) = 1/C(t+1, k)입니다.

처음 k개를 넣은 직후에는 가능한 조합이 하나뿐이라 가정이 성립하고, 한 단계씩 진행해도 유지됩니다. 따라서 끝까지 읽었을 때 C(n, k)가지 조합이 모두 같은 확률입니다.

흔한 실수와 경계 조건

범위를 하나 틀리는 실수. 0부터 세는 반복문에서 i는 (i+1)번째 원소입니다. 이때 randomBelow(i + 1)이 아니라 randomBelow(i)를 쓰면 들어갈 확률이 k/(i+1)이 아니라 k/i가 됩니다. 1 ≤ k < n이라고 하면, i = k, 곧 k+1번째 원소에서는 randomBelow(k)가 항상 k보다 작은 값을 주므로 이 원소가 무조건 들어갑니다. 같은 가정에서 앞 절과 같은 방식으로 끝까지 계산하면(이 글의 계산) 처음 k개가 남을 확률은 (k−1)/(n−1), 나머지 원소는 k/(n−1)이 됩니다. 처음 k개만 불리해지는 치우친 표본입니다. k = 0이면 이 잘못된 코드는 첫 원소에서 randomBelow(0)을 부릅니다. 주석에 적은 범위(1 ≤ n) 밖이므로 이 경우는 확률을 따질 것도 없이 잘못된 호출입니다. n ≤ k이면 잘못된 호출까지 가지 않아 위 확률식도 적용되지 않습니다. 위 코드는 seen을 먼저 올리고 randomBelow(seen)을 부르는 식으로 1부터 센 개수를 그대로 써서 이 혼동을 줄였습니다.

입력이 k개보다 적을 때. 위 함수는 들어온 원소를 모두 담아 k개보다 짧은 배열을 돌려줍니다. 길이를 모르는 스트림이라 이 경우를 미리 막을 수 없으니, 오류로 볼지 짧은 결과를 받아들일지는 호출하는 쪽에서 정해야 합니다. 참고로 Python의 random.sample(population, k)는 k가 모집단보다 크면 ValueError를 냅니다.

k가 0 이하일 때. 위 코드는 k = 0이면 입력을 읽지 않고 빈 배열을 돌려주고, 음수나 정수가 아닌 값은 RangeError로 거부합니다. 다른 정책도 가능하지만 어느 쪽이든 명시해 두는 편이 낫습니다. 음수 k를 그대로 두면 seen <= k가 한 번도 참이 되지 않아 조용히 빈 결과가 나오기 쉽습니다.

결과의 순서. 저수지에 담긴 원소의 집합은 고르지만 배열 안의 순서까지 무작위인 것은 아닙니다. 입력이 k개 이하면 입력 순서 그대로이고, 처음 k개 중 끝까지 남은 원소는 원래 자리를 지킵니다. 순서가 의미 있는 곳에 쓴다면 결과를 Fisher–Yates로 한 번 더 섞으세요.

난수의 품질. 위 증명은 M이 1부터 j까지 정확히 고르게 나온다는 전제 위에 있습니다. 실수 난수에 j를 곱해 버림하는 방식은 j가 커질수록 정밀도 한계를 따져 봐야 합니다. Python 문서는 기본 생성기인 메르센 트위스터가 결정적이어서 암호 용도에는 전혀 맞지 않고, 보안 용도에는 secrets 모듈을 쓰라고 안내합니다. 표본을 누가 예측하면 곤란한 상황이라면 알고리즘보다 난수원부터 따져야 합니다.

길이를 알 때, 더 빨라야 할 때

모든 원소가 이미 메모리에 있고 개수도 안다면 저수지 샘플링이 필요 없습니다. Fisher–Yates 셔플을 k번만 돌리거나, Python이라면 random.sample을 쓰면 됩니다. 다만 random.sample은 모집단으로 시퀀스를 받습니다. Python 3.11부터는 집합도 자동으로 리스트로 바꿔 주지 않습니다. 스트림에 쓰려면 결국 전부 리스트로 만들어야 하니, 한 번 읽고 끝내는 상황에는 맞지 않습니다.

반대로 n이 k보다 훨씬 크면 대부분의 원소는 난수를 하나 뽑은 뒤 버려집니다. DADS는 Jeffrey Scott Vitter의 1985년 논문 “Random Sampling with a Reservoir”(ACM TOMS 11(1))를 소개하면서, 그 논문의 Algorithm X, Y, Z는 원소마다 건너뛸지 정하는 대신 몇 개를 건너뛸지를 한 번에 정해서 난수를 훨씬 적게 쓴다고 설명합니다. 난수 생성 비용이 문제가 될 때 찾아볼 만한 방향입니다.

참고 자료

모두 2026-10-13에 확인했습니다.

블로그로