linked list
random
math
]
Leetcode 0382. Linked List Random Node
https://leetcode.com/problems/linked-list-random-node
Let us solve this question for follow-up question: we do not want to use additional memory here. There is specific method for this, whith is called reservoir sampling (actually, special case of it), which I am going to explain now. Imagine, that we have only 3
nodes in our linked list, then we do the following logic:
- On the first step we either stay in first node with probability
1/2
or go to the next one also with probability1/2
. - On the second step, we do the following: if we are in node number 1, then we stay in this node with probability
2/3
and go to node number 2 with probability1/3
. If we in node number 2, then we stay in it with probability1/3
and go to node number 3 with probability2/3
. Why everything will be OK? There is only one way, how you can stay in node 1, and probability is1/2 * 2/3
. There are two cases, how you can go to node number 2: stay on first step and go on second and go on first step and stay on second: probability will be1/2*1/3 + 1/2*1/3 = 1/3
. Finally, probability to be in node 3 is also1/3
.
Now, imagine, that we already covered n
nodes and we have probabilites: 1/n, 1/n, ..., 1/n
. Then we need to make one more step and have probabilites 1/(n+1), 1/(n+1), ... , 1/(n+1)
. How we can do it?
-
If we are in the first node, we stay in it with probability
n/(n+1)
and go next with probability1/(n+1)
. -
If we are in the second node, we stay in it with probability
(n-1)/(n+1)
and go next with probability2/(n+1)
.n. If we are in node
n
, we stai in it with probability1/(n+1)
and go next with probabilityn/(n+1)
.
Why it is working? There is only one way how you can be in node 1: with probability 1/n* n/(n+1) = 1/(n+1)
. There is two options for node 2, with probability 1/n*(n-1)/(n+1) + 1/n* 1/(n+1) = 1/(n+1)
and so on.
Finally, let us go to the code: we keep n
and k
values: n
is number of nodes we count in list so far and k
is current node we are in. Each moment of time we decide if we go to next node or rest in current. If we go to next, we increase k
by one and change our ans
node. Increase n
in any case.
Complexity: time compexity is O(n)
and you can not really do anything with if you are not allowed to use extra memory. Space complexity however only O(1)
, if we do not count our input.
class Solution:
def __init__(self, head):
self.head = head
def getRandom(self):
n, k = 1, 1
head, ans = self.head, self.head
while head.next:
n += 1
head = head.next
if random.random() < k/n:
ans = ans.next
k += 1
return ans.val
If you like the solution, you can upvote it on leetcode discussion section: Problem 0382