描述
给定一个单链表,随机选择链表的一个节点,并返回相应的节点值。保证每个节点被选的概率一样。
进阶:
如果链表十分大且长度未知,如何解决这个问题?你能否使用常数级空间复杂度实现?
示例:
// 初始化一个单链表 [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();
解析
先说结论,当你遇到第 i
个元素时,应该有 1/i
的概率选择该元素,1 - 1/i
的概率保持原有的选择。
举例说明: 1 - 10
遇到1,概率为1,保留第一个数。
遇到2,概率为1/2,这个时候,1和2各1/2的概率被保留
遇到3,3被保留的概率为1/3,(之前剩下的数假设1被保留),2/3的概率 1 被保留,(此时1被保留的总概率为 2/3 * 1/2 = 1/3)
遇到4,4被保留的概率为1/4,(之前剩下的数假设1被保留),3/4的概率 1 被保留,(此时1被保留的总概率为 3/4 * 2/3 * 1/2 = 1/4)
以此类推,每个数被保留的概率都是1/N。
代码
/* 返回链表中一个随机节点的值 */ int getRandom(ListNode head) { Random r = new Random(); int i = 0, res = 0; ListNode p = head; // while 循环遍历链表 while (p != null) { // 生成一个 [0, i) 之间的整数 // 这个整数等于 0 的概率就是 1/i if (r.nextInt(++i) == 0) { res = p.val; } p = p.next; } return res; }
升级--随机选择k个节点
如果要随机选择 k
个数,只要在第 i
个元素处以 k/i
的概率选择该元素,以 1 - k/i
的概率保持原有选择即可。
/* 返回链表中 k 个随机节点的值 */ int[] getRandom(ListNode head, int k) { Random r = new Random(); int[] res = new int[k]; ListNode p = head; // 前 k 个元素先默认选上 for (int j = 0; j < k && p != null; j++) { res[j] = p.val; p = p.next; } int i = k; // while 循环遍历链表 while (p != null) { // 生成一个 [0, i) 之间的整数 int j = r.nextInt(++i); // 这个整数小于 k 的概率就是 k/i if (j < k) { res[j] = p.val; } p = p.next; } return res; }