382. 链表随机节点

题目描述

输入一个单链表，链表用数组表示，例如 [1,2,3]。实现类 Solution：

Solution(head)：用链表头节点 head 创建对象。getRandom()：随机挑一个节点，返回它的值。每个节点被挑中的概率必须相等。

示例中第一行为方法名，第二行为参数，构造函数的返回记为 null。

示例 1

输入：操作 = ["Solution","getRandom","getRandom","getRandom","getRandom","getRandom"]
参数 = [[[1,2,3]],[],[],[],[],[]]
输出：[null, 3, 1, 3, 2, 2]

解释：每次 getRandom() 都在 1、2、3 中等概率选一个，这次运行依次得到 3,1,3,2,2，结果每次运行可能不同。

数据范围

链表节点数在 [1, 10^4] 范围内
-10^4 <= Node.val <= 10^4
getRandom 最多调用 10^4 次

进阶：如果链表极长且事先不知道长度，怎么办？能否不使用额外空间？

解题思路：水塘抽样

水塘抽样：从头遍历链表，走到第 i 个节点时，以 1/i 的概率用它替换当前答案（即 randrange(i) == 0），否则保留旧答案。遍历结束时的答案就是等概率抽出的节点。不需要提前知道长度，也不需要把链表拷到数组里。

为什么正确

处理完前 i 个节点后，ans 是这 i 个节点中的每一个的概率都是 1/i。第 i+1 个节点以 1/(i+1) 替换它，旧节点留下的概率变成 1/i * i/(i+1) = 1/(i+1)。

复杂度分析

构造 O(1)；每次 getRandom 时间 O(n)，额外空间 O(1)。

易错点

替换概率是 1/i，不是固定的 1/2：如果每次都招硬币，越靠后的节点越容易被选中。另外 i 从 1 开始计数，randrange(0) 会报错。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def __init__(self, head: Optional[ListNode] = None):
        self.head = head

    def getRandom(self) -> int:
        ans, i, cur = 0, 0, self.head
        while cur:
            i += 1
            if random.randrange(i) == 0:
                ans = cur.val
            cur = cur.next
        return ans

class Solution {
    private ListNode head;
    private java.util.Random rnd = new java.util.Random();
    public Solution(ListNode head) {
        this.head = head;
    }

    public int getRandom() {
        int ans = 0, i = 0; ListNode cur = head;
        while (cur != null) {
            i++;
            if (rnd.nextInt(i) == 0)
                ans = cur.val;
            cur = cur.next;
        }
        return ans;
    }
}

class Solution {
public:
    ListNode* head;
    Solution(ListNode* head) {
        this->head = head;
    }

    int getRandom() {
        int ans = 0, i = 0; ListNode* cur = head;
        while (cur) {
            i++;
            if (rand() % i == 0)
                ans = cur->val;
            cur = cur->next;
        }
        return ans;
    }
};

type Solution struct {
	head *ListNode
}

func Constructor(head *ListNode) Solution {
	return Solution{head}
}

func (this *Solution) GetRandom() int {
	ans, i, cur := 0, 0, this.head
	for cur != nil {
		i++
		if rand.Intn(i) == 0 {
			ans = cur.Val
		}
		cur = cur.Next
	}
	return ans
}

var Solution = function(head) {
    this.head = head;
};

Solution.prototype.getRandom = function() {
    let ans = 0, i = 0, cur = this.head;
    while (cur) {
        i++;
        if (Math.floor(Math.random() * i) === 0)
            ans = cur.val;
        cur = cur.next;
    }
    return ans;
};

class Solution {
    head: ListNode | null;
    constructor(head: ListNode | null) {
        this.head = head;
    }

    getRandom(): number {
        let ans = 0, i = 0, cur = this.head;
        while (cur) {
            i++;
            if (Math.floor(Math.random() * i) === 0)
                ans = cur.val;
            cur = cur.next;
        }
        return ans;
    }
}

载入编辑器，自己跑一遍 →

其他解法

记录数组用 O(n) 空间换取 O(1) 的 getRandom；水塘抽样每次 O(n)，但只需 O(1) 空间，也适用于长度未知的数据流。

解法 2：记录所有节点值时间 O(n)，空间 O(n)

构造时遍历一次链表，把所有节点值存进数组。getRandom 时生成 [0, len) 内的随机下标，返回对应的值。构造 O(n)，每次 getRandom O(1)。

调用 getRandom 次数多时更快。需要 O(n) 额外空间，并且要求链表长度已知、不再变化；链表极长或是数据流时只能用水塘抽样。

class Solution:
    def __init__(self, head: Optional[ListNode] = None):
        self.arr = []
        while head:
            self.arr.append(head.val)
            head = head.next

    def getRandom(self) -> int:
        return self.arr[random.randrange(len(self.arr))]

class Solution {
    private List<Integer> arr = new ArrayList<>();
    private java.util.Random rnd = new java.util.Random();

    public Solution(ListNode head) {
        while (head != null) {
            arr.add(head.val);
            head = head.next;
        }
    }

    public int getRandom() {
        return arr.get(rnd.nextInt(arr.size()));
    }
}

class Solution {
    vector<int> arr;
public:
    Solution(ListNode* head) {
        while (head) {
            arr.push_back(head->val);
            head = head->next;
        }
    }

    int getRandom() {
        return arr[rand() % arr.size()];
    }
};

type Solution struct {
	arr []int
}

func Constructor(head *ListNode) Solution {
	arr := []int{}
	for ; head != nil; head = head.Next {
		arr = append(arr, head.Val)
	}
	return Solution{arr}
}

func (this *Solution) GetRandom() int {
	return this.arr[rand.Intn(len(this.arr))]
}

var Solution = function(head) {
    this.arr = [];
    for (; head; head = head.next) this.arr.push(head.val);
};

Solution.prototype.getRandom = function() {
    return this.arr[Math.floor(Math.random() * this.arr.length)];
};

class Solution {
    arr: number[] = [];
    constructor(head: ListNode | null) {
        for (; head; head = head.next) this.arr.push(head.val);
    }

    getRandom(): number {
        return this.arr[Math.floor(Math.random() * this.arr.length)];
    }
}

链表专题的其他题目

查看链表全部题目 →