LC 0398 [M] Random Pick Index - ALawliet/algorithms GitHub Wiki

in reservoir sampling, we want to maintain the chance of selecting other vs. selecting 0 as the sample size grows

if we have 1,2,3,3,3, select 3
i=2, reservoir=[0], we have a 1/1 chance to select 0
i=3, reservoir=[0,1], we have a 1/1 * 1/2 = 1/2 chance to select other, 1/2 chance to select 0
i=4, reservoir=[0,1,2], we have a 1/1 * 1/2 * 2/3 = 1/3 chance to select other, a 1/3 chance to select 0

think about it as if the target is 3 and the sample is all 3s, maintain the probability as we discover targets (in this case, all of them) and it grows

0 1 2 (indices)
3 3 3 (values)

step 1: select the first sample "3" (index 0), the probability to select this sample is 1/1 and the probability that this selected item will remain in the reservior is 1/3 (other 2 samples could replace it), setting the resulting probablity to 1/1 * 1/3 -> 1/3

step 2: select out of the first two samples "3 3" (index 0, index 1), the probability to select the second sample is 1/2 and the probability that it will remain in the reservior is 2/3 (other 1 sample could replace it), setting the resulting probablity to 1/2 * 2/3 -> 1/3

step 3: select out of the first three samples "3 3 3" (index 0, index 1, index 2), the probability to select the third sample is 1/3 and the probability that it will remain in the reservior is 3/3 or simply 1.0 (none other samples could replace it), setting the resulting probablity to 1/3 * 1 -> 1/3

class Solution:

    def __init__(self, nums):
        self.nums = nums

    def pick(self, target):
        index = None
        chance = 0 # chance to occur
        for i in range(len(self.nums)):
            if self.nums[i] == target: # from the target sample size
                chance += 1 # with possible duplicates 
                roll = random.randint(1, chance) # maintain chance as sample size grows
                if roll == chance:
                    index = i
        return index


# Your Solution object will be instantiated and called as such:
# obj = Solution(nums)
# param_1 = obj.pick(target)

class Solution:

    def __init__(self, nums):
        self.nums = nums

    def pick(self, target):
        res = None
        k = 0
        for i, x in enumerate(self.nums):
            if x == target:
                k += 1
                if random.randint(1, k) == k:
                    res = i
        return res