题目描述

[EN | CN]

给定一个单链表,随机选择链表的一个结点,并返回相应的结点值。保证每个结点 被选的概率一样

进阶:

如果链表十分大且长度未知,如何解决这个问题?你能否使用常数级空间复杂度实现?

示例:

1
2
3
4
5
6
7
8
// 初始化一个单链表 [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()方法应随机返回1,2,3中的一个,保证每个元素被返回的概率相等。
solution.getRandom();

解法 1:蓄水池采样

该题目的本质其实是 数据流中的随机抽样问题:从长度未知的链表中随机取 $k$ 个元素,保证每个元素被取到的概率相同。

对于这种问题,一种经典的解法是 蓄水池采样算法(Reservoir Sampling),算法步骤如下:

  • 对于前 $k$ 个元素,全部保留。
  • 对于第 $i$ 个元素($i > k$),以 $\frac{k}{i}$ 的概率保留它,以 $\frac{i - k}{i}$ 的概率丢弃它。

用数学归纳法证明如下:

对于前 $k$ 次采样,前 $k$ 个元素都被保留,每个元素被取到的概率都为 $1$。

对于第 $i = k + 1$ 次采样,以 $\frac{k}{i} = \frac{k}{k + 1}$ 的概率保留,之前取到的 $k$ 个元素被取到的概率变为:

$$ 1 \times (1 - \frac{k}{k + 1} \cdot \frac{1}{k}) = \frac{k}{k + 1} $$

对于第 $i = k + 2$ 次采样,以 $\frac{k}{i} = \frac{k}{k + 2}$ 的概率保留,之前取到的 $k$ 个元素被取到的概率变为: $$ \frac{k}{k + 1} \times (1 - \frac{k}{k + 2} \cdot \frac{1}{k}) = \frac{k}{k + 2} $$

……

对于第 $i$ 次采样($k > k$),以 $\frac{k}{i}$ 的概率保留,之前取到的 $k$ 个元素被取到的概率变为:

$$ \frac{k}{i - 1} \times (1 - \frac{k}{i} \cdot \frac{1}{k}) = \frac{k}{i} $$ 即证。

假设链表长度为 $n$,那么

  • 时间复杂度为 $O(n)$,只要遍历到 $n$ 即可;
  • 空间复杂度为 $O(1)$,不需要额外空间。

实现与结果如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import random

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

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, ret, count = self.head.next, self.head.val, 1

        while node:
            count += 1
            r = random.random()
            if r < 1 / count:
                ret = node.val
            node = node.next

        return ret

# Your Solution object will be instantiated and called as such:
# obj = Solution(head)
# param_1 = obj.getRandom()
  • 执行用时:168 ms,在所有 Python3 提交中击败了 61.98% 的用户。
  • 内存消耗:17.1 MB,在所有 Python3 提交中击败了 50.00% 的用户。