398. Random Pick Index

leetcode.com/problems/random-pick-index/

Random Pick Index - LeetCode

class Solution:

    def __init__(self, nums: List[int]):
        self.nums = nums
        
    def pick(self, target: int) -> int:
        
        output = 0
        count = 0
        
        for i, x in enumerate(self.nums) :
            
            if x != target :
                continue
            
            count += 1
            chance = random.randint(1, count)
            
            if chance == count :
                output = i
                    
        return output
                

 

정리 : Reservoir Sampling 사용하면 랜덤 인덱스를 O(N)으로 구할 수 있음.

ex) 만약, [1,2,3,3,3] 이라는 배열이 있다고 가정해보자.

i=2가 나올 확률 = 1*1/2*2/3

i=3가 나올 확률 = 1/2*2/3

i=4가 나올 확률 = 1/3