Select a random element from the stream of elements

//This is GFG code,You won't get any practice for this problem.
'''
Logic is simple: We are passing the stream of numbers into the selectRandom function.
We maintain a count of numbers seen so far in a variable `count`.
For each new number, we generate a random number from 0 to count-1.
If the random number equals count-1, we update the probable element to the current number;
otherwise, we keep the previous result. This ensures each element has an equal probability
of being selected from the stream, even if the total number of elements is unknown in advance.

Example: [1, 2, 3, 4]

- First element (1): Only one element exists, so its probability of being chosen is 100%.
  `probable_element = 1`.

- Second element (2): Now we have two elements [1, 2].
  Each element should have a 50% chance of being selected.
  Generating a random number between 0 and 1 gives each element equal probability.
  If the random number is 1, the current number (2) becomes the `probable_element`;
  otherwise, it remains 1.

- Third element (3): Now we have three elements [1, 2, 3].
  Each element should have a 1/3 chance of being chosen.
  Generating a random number between 0 and 2 achieves this.
  If the number is 2, we update `probable_element` to 3; otherwise, it remains the previous choice.

- Fourth element (4): With four elements [1, 2, 3, 4], each has a 1/4 chance.
  A random number between 0 and 3 is generated.
  If the number is 3, we update `probable_element` to 4; otherwise, we keep the previous element.

This approach is known as **reservoir sampling**, allowing us to select a random element from a stream
of unknown size with uniform probability.

Note: the n can be anything but at every ith point, all the elements got the equal probabilty.
'''
import random

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

    def pick(self, target):
        res = -1   # current chosen index
        count = 0  # number of target indices seen so far

        for i, num in enumerate(self.nums):
            if num == target:
                count += 1
                # replace current result with probability 1/count
                if random.randint(0, count-1) == count-1:
                    res = i
        return res

# Example
nums = [1, 2, 3, 3, 3]
solution = Solution(nums)
print(solution.pick(3))  # Randomly returns 2, 3, or 4
print(solution.pick(1))  # Returns 0
print(solution.pick(3))  # Randomly returns 2, 3, or 4