528. Random Pick with Weight (Array, Math, Binary Search, Prefix Sum, Randomized)

링크 - https://leetcode.com/problems/random-pick-with-weight/description/?envType=company&envId=facebook&favoriteSlug=facebook-thirty-days

난이도 - Medium

Intuition

이 문제는 Prefix Sums와 Binary Search를 통합해 풀 수 있다. Linear Search로도 풀 수는 있지만, pickIndex()에서 더 나은 시간복잡도를 가지려면 Binary Search를 활용할 수 있다.

1. 먼저 w에 있는 값들을 더하면서 self.prefix_sums에 append한다. 끝나면 다 더한 값을 self.total_sum에 저장한다.

2. self.total_sum * random.random()을 통해 0 - self.total_sum안에 있는 랜덤한 숫자를 선택한다.

3. left, right를 0, len(self.prefix_sums) - 1로 설정한다.

4. left와 right의 중간값 mid를 구한뒤 self.prefix_sums[mid]값을 구한다. 만약 찾고자 하는 값이 이거보다 크면 left를 mid+1로 하고, 아니면 right를 mid로 한다.

5. 4를 반복하고 left를 반환한다.

Code

class Solution:

    def __init__(self, w: List[int]):
        self.prefix_sums = []
        prefix_sum = 0
        for weight in w:
            prefix_sum += weight
            self.prefix_sums.append(prefix_sum)
        self.total_sum = prefix_sum

    def pickIndex(self) -> int:
        target = self.total_sum * random.random()
        left, right = 0, len(self.prefix_sums) - 1
        while left < right:
            mid = left + (right - left) // 2
            if target > self.prefix_sums[mid]:
                left = mid + 1
            else:
                right = mid

        return left


# Your Solution object will be instantiated and called as such:
# obj = Solution(w)
# param_1 = obj.pickIndex()

Complexity

Time Complexity: __init__ - O(n), pickIndex() - O(logN)

Space Complexity: __init__ - O(n), pickIndex() - O(1)