398. Random Pick Index

class Solution:
    def __init__(self, nums: List[int]):
        self.A = nums
    def pick(self, target: int) -> int:
        """
        tc O(N)  sc O(1) 
        use cnt to record by index i, the amount of number target
        if we choose first number as sample position, by array[:i+1] the probability of choosing first target number is 
        1 * (1-1/2) * (1-1/3) * ... (1-1/n)  = 1/n   
        Hence, by the time array iteration ends, the last time that we get random number at first position is universal probability  
        """
        n = len(self.A)
        cnt = idx = 0
        for i in range(n): 
            if self.A[i] == target: #  
                cnt += 1
                #random.randint(0,cnt-1) < k can be used for k/n probability
                if random.randint(0,cnt-1) == 0:     
                    idx = i 
        return idx 

class Solution:

    def __init__(self, nums: List[int]):
        """
        main idea:  hash map  
        tc O(N)  sc O(N)
        """
        self.d = collections.defaultdict(list)
        for i, x in enumerate(nums):
            self.d[x].append(i)

    def pick(self, target: int) -> int:
        """
        tc O(1) sc O(1)
        """
        lst = self.d[target]
        i = random.randint(0,len(lst)-1)
        return lst[i]