398. Random Pick Index

solution

# 蓄水池算法: 给定数据流，长度N很大，且N直到处理完所有数据之前都不可知，如何在只遍历一遍数据（O(N)）的情况下，随机选取出m个不重复的数据
# i <= m: i从1开始，小于m则直接入池
# i > m: 从[1, i]中选一个数x, x > m 掠过；x<=m, x位置的元素替换为i位置的元素
# https://zhuanlan.zhihu.com/p/119329875

class Solution:
    def __init__(self, nums: List[int]):
        self.nums = nums

    def pick(self, target: int) -> int:
        count = 0
        res = None
        for i, num in enumerate(self.nums):
            if num == target:
                count += 1
                if random.randint(1, count) == 1:
                    res = i
        return res

class Solution:
    def __init__(self, nums: List[int]):
        self.nums = nums

    def pick(self, target: int) -> int:
        res = None
        count = 0

        for i, num in enumerate(self.nums):
            if num == target:
                if res is None:
                    res = i  # 只选一个, 第一个先入池
                else:
                    if random.randint(0, count) == 0:  # 大于1后, 替代之前的概率 1/(cnt-1) 取决于count是否加上了自身
                        res = i

                count += 1
        return res

时间复杂度：O(n) 空间复杂度：O(1)

class Solution:
    def __init__(self, nums: List[int]):
        self.indices = defaultdict(list)
        for i, num in enumerate(nums):
            self.indices[num].append(i)

    def pick(self, target: int) -> int:
        return random.choice(self.indices[target])

follow up-Reservoir Sampling/Rejection Sampling类

多线程版本

# 有权重(矩阵面积)抽取矩阵
import random

class Solution:
    def __init__(self, rects: List[List[int]]):
        self.rects = rects
        self.areas = []
        self.total_area = 0

        for rect in rects:
            area = (rect[2] - rect[0] + 1) * (rect[3] - rect[1] + 1)
            self.total_area += area
            self.areas.append(self.total_area)

    def pick(self) -> List[int]:
        r = random.randint(1, self.total_area)
        for i in range(len(self.rects)):
            if r <= self.areas[i]:
                rect = self.rects[i]
                x = random.randint(rect[0], rect[2])
                y = random.randint(rect[1], rect[3])
                return [x, y]

# 从左往右扫描，每次随机选择当前位置或当前位置右侧的某一个数与当前位置进行交换，相当于随机抽取一个

class Solution:
    def __init__(self, nums: List[int]):
        self.nums = nums
        self.original = copy.deepcopy(nums)

    def reset(self) -> List[int]:
        self.nums = copy.deepcopy(self.original)
        return self.nums

    def shuffle(self) -> List[int]:
        n = len(self.nums)
        for i in range(n):
            j = random.randrange(i, n)
            self.nums[i], self.nums[j] = self.nums[j], self.nums[i]
        return self.nums

# 拒绝采样 Rejection Sampling: 对于当前随机生成器的一部分生成情况进行丢弃，从而完成实现利用给定的随机生成器得到我们需要的概率分布
# 如何用骰子得到1/7概率？把骰子扔两次，获得 6 * 6 = 36 个样本，丢弃最后一个样本，剩下的 35 个样本平分成 7 份，对应的概率值便为 1/7

class Solution:
    def rand10(self) -> int:
        while True:
            r1 = rand7()
            r2 = rand7()
            idx = (r1 - 1) * 7 + r2  # 减1 变成[0,6] 的随机数，乘7 变为{0,7,14,21,28,35,42} 中的随机数，加col 变为[1,49]的随机数
            if idx <= 40:  # 拒绝一部分，只保留[1, 40]
                return 1 + idx % 10  # [1, 10]

# general solution: creating randM() using randN()
class Solution:
    def randM(self, N, M):
        # acceptable is the desired range which can generate required integer directly
        curr = acceptable = N * N - (N * N) % M
        # if current no is not in the acceptable range, discard it and repeat the process again
        while curr >= acceptable:
            curr = (randN() - 1) * N + randN() - 1
        return curr % M + 1

# 拒绝采用: 随机角度 + 随机半径中一点

class Solution:
    def __init__(self, radius: float, x_center: float, y_center: float):
        self.radius = radius
        self.x_center = x_center
        self.y_center = y_center

    def randPoint(self) -> List[float]:
        degree = random.uniform(0, 1) * 2 * math.pi
        length = math.sqrt(random.uniform(0, 1)) * self.radius  # 注意 sqrt, pi * r^2
        x = self.x_center + length * math.cos(degree)
        y = self.y_center + length * math.sin(degree)
        return [x, y]

# 第一次遇到最大的，就把res update 成当前index
# 第二次遇到最大的，就1/2概率把res update 成当前index
# 第三次遇到最大的，就1/3概率把res update 成当前index

def max_random_index(nums):
    max_val = float('-inf')
    max_index = -1
    count = 0

    for i, num in enumerate(nums):
        if num > max_val:
            max_val = num
            max_index = i
            count = 1
        elif num == max_val:
            count += 1

            # Probability of 1/count
            if random.randint(0, count - 1) == 0:
                max_index = i

        print(max_index, end=" ")

## follow up， very long array can not fit into memory

import random

def sample_gen(n, forbid):
    state = dict()
    track = dict()
    for (i, o) in enumerate(forbid):
        x = track.get(o, o)
        t = state.get(n-i-1, n-i-1)
        state[x] = t
        track[t] = x
        state.pop(n-i-1, None)
        track.pop(o, None)
    del track
    for remaining in range(n-len(forbid), 0, -1):
        i = random.randrange(remaining)
        yield state.get(i, i)
        state[i] = state.get(remaining - 1, remaining - 1)
        state.pop(remaining - 1, None)

Previous382. Linked List Random Node Next29. Divide Two Integers

Last updated 27 days ago