Leetcode: Linked List Random Node (Reservoir Sampling: 저수지 샘플링)

Given a singly linked list, return a random node's value from the linked list. Each node must have the same probability of being chosen.

Follow up:
What if the linked list is extremely large and its length is unknown to you? Could you solve this efficiently without using extra space?

Example:

// Init a singly linked list [1,2,3].
ListNode head = new ListNode(1);
head.next = new ListNode(2);
head.next.next = new ListNode(3);
Solution solution = new Solution(head);

// getRandom() should return either 1, 2, or 3 randomly. Each element should have equal probability of returning.
solution.getRandom();

 

Linked List에서 random하게 element를 뽑는 방법을 구하는 문제이다. December challenge이 Day2문제여서 접하게 되었다. 간단하게 풀자면 ListNode를 끝까지 순회하며 array에 저장해두고, getRandom()에서 0, len(arr)-1의 range에서 random한 숫자를 뽑고 배열의 해당 인덱스에 있는 값을 리턴하면 된다.

Naive Solution

class Solution:

    def __init__(self, head: ListNode):
        """
        @param head The linked list's head.
        Note that the head is guaranteed to be not null, so it contains at least one node.
        """
        self.arr = []
        idx, node = 0, head
        while node:
            self.arr.append(node.val)
            idx, node = idx+1, node.next
        self.length = idx

    def getRandom(self) -> int:
        """
        Returns a random node's value.
        """
        rand = random.randint(0, self.length-1)
        return self.arr[rand]

위와 같이 해결하면, getRandom()은 Space Complexity O(N)이 된다. 다만 이 문제의 설명 끝에, list의 길이가 길고, 그 길이가 알려져있지 않은 경우를 가정했을 때, extra space를 사용하지 않고 풀어보라는 코멘트가 나와서 끙끙 앓다가, 결국 리트코드에서 제공해주는 답지를 봤다. 답지에서는 Reservoir Sampling (저수지 샘플링)이라는 방법을 사용한다. 

 

Reservoir Sampling

Reservoir Sampling은 n개의 list에서 k개의 random element를 equal probablility로 뽑는 방법이다. 결론부터 말하면 코드는 아래와 같다.

# S has items to sample, R will contain the result
def ReservoirSample(S[1..n], R[1..k])
  # fill the reservoir array
  for i := 1 to k
      R[i] := S[i]

  # replace elements with gradually decreasing probability
  for i := k+1 to n
    # randomInteger(a, b) generates a uniform integer
    #   from the inclusive range {a, ..., b} *)
    j := randomInteger(1, i)
    if j <= k
        R[j] := S[i]

R이 random sampling된 결과를 가진 array라고 할때, 아래와 같은 과정을 거친다.

 

1. R[1..k]는 [1..n]과 같다

2. k+1부터 n까지의 i에 대하여, random하게 i 이하의 수를 뽑는다. 만약 이 수(j)가 k보다 작거나 같다면, R[j]를 S[i]로 교체한다.

 

즉 time complexity는 O(N), space complexity는 O(k)이다. 

이제 이 방법이 어떻게 equal probability를 가지는지 이해하는 것이 관건이다. 

 

1번 과정에서, 1..n에 있는 i에 대해서 1의 확률로 R에 포함이 된다. 하지만, 2번 과정에서, random하게 1에서 포함되었던 값이 교체될 가능성이 있다. k이상의 i값에 대하여, 1과 i 사이의 random한 j가 k 이하일 경우 교체가 이루어지므로, S[1..i]의 각 element가 교체되어 사용되지 않을 확률은 1/i 이다. 이 말은, 교체가 되지 않을 확률은 1 - 1/i = (i-1) / i 라는 것이다. 

이를 토대로 정리했을 때, 각 element가 R에 들어갈 확률은 아래와 같이 계산할 수 있다.

 

S[i] for i in [1..k]

$$ p = 1 * \frac{k}{k+1} * \frac{k+1}{k+2} * ... * \frac{n-2}{n-1} * \frac{n-1}{n} = \frac{k}{n}$$

1..k 사이의 i에 대해서는, 위의 1번 과정에서 반드시 R에 포함이 된다. 즉, 1의 확률을 가지고, 이후의 2번 과정에서 다른 element로 교체가 되지 않아야 한다. k+1..n 사이의 인덱스 x에서 다른 값으로 교체되지 않을 확률은 위에서 살펴보았듯이 x-1 / x 이므로, 이를 k+1 과 n사이이 모든 인덱스에 대해 계산하면 k/n 이 된다.

 

S[i] for i in [k+1..n]

$$ p = \frac{k}{i} * \frac{i}{i+1} * \frac{i+1}{i+2} * ... * \frac{n-2}{n-1} * \frac{n-1}{n} = \frac{k}{n} $$

k+1..n 사이의 i에 대해서는, 2번 과정에서 i번 인덱스에 대해 순회를 할 때에 확률적으로 R에 포함이 될 수 있다. 우리는 S[i]가 포함될 확률에 대해서만 알고 싶으므로, i가 포함될 가능성이 없는 1..i-1번 인덱스에서는 어떠한 일이 일어나 어떠한 원소가 포함되든 알 필요가 없다. S[i]가 i번째에 포함이 되더라도, 최종적으로 R에 남아있으려면 이후 인덱스에서 다른 값으로 교체되지 않아야 한다. 즉, 최초 i번째 인덱스에서 포함이 될 확률 k/i에, 이후로부터 마지막 인덱스 n까지 다른 값으로 교체되지 않을 확률들을 모두 곱해서 계산하면 k/n이 된다. 

 

이처럼, 1부터 n까지의 모든 인덱스에 대하여 S[i]가 R에 포함될 확률은 k/n으로 계산된다.

이를 바탕으로, 주어진 문제를 푼 코드는 아래와 같다.

 

Solution Using Reservoir Sampling

class Solution:

    def __init__(self, head: ListNode):
        """
        @param head The linked list's head.
        Note that the head is guaranteed to be not null, so it contains at least one node.
        """
        self.head = head

    def getRandom(self) -> int:
        """
        Returns a random node's value.
        """
        node = self.head
        scope = 1
        answer = node.val
        
        while node:
            j = random.randint(1, scope)
            if j <= 1:
                answer = node.val
            scope += 1
            node = node.next
        return answer

이 문제는 k=1인 Reservoir sampling문제이므로, 위의 알고리즘에서 k값만 1로 교체한 거의 동일한 코드로 정답을 찾을 수 있다.

노드를 1부터 n까지 모두 순회해야 하기 때문에 time complexity는 O(n)이고, space complexity는 O(1)이다. (왜 노드를 처음부터 끝까지 순회해야하는지가 이해가 가지 않으면 위의 확률을 계산하는 부분을 다시 한 번 참고하면 좋을 것 같다.)

 

(이 포스팅은 Leetcode의 답지를 참고하여 제가 이해한대로 정리한 글이므로, 확률 계산과 같은 부분은 다소 부정확한 부분이 있을 수 있습니다.)

 

 

링크: leetcode.com/problems/linked-list-random-node/