Random Pick With Weight

0528.Random-Pick-with-Weight/memo.md

@@ -0,0 +1,34 @@
+# 528. Random Pick with Weight
+
+## step1
+9mぐらいでまず naive を書く。計算量 O(N)。
+
+bisectを使って高速化。計算量 O(log N)
+
+## step2
+いい加減な名前で書いてしまったので改善
+
+乱数の書き方を調べる:
+
+- random.randint: uniformと同じ書き方で書ける。もともと整数なのでこちらの方が良さそう
+- random.choices: `return random.choices(list(range(len(self.prefix_sums))), cum_weights=self.prefix_sums)[0]`
+    - 内部で二分探索が走っている
+    - https://github.com/python/cpython/blob/219768ff531fc0686de623139562ee9f9537df98/Lib/random.py#L460
+- np.random.choice: O(N)だがバッチジョブだと高速化
+
+---
+
+Alias method
+
+https://leetcode.com/problems/random-pick-with-weight/solutions/671439/python-smart-o1-solution-with-detailed-e-r0gx/?envType=problem-list-v2&envId=7p55wqm
+
+https://en.wikipedia.org/wiki/Alias_method
+
+O(N)の前計算を行なっておくことでO(1)で生成できる
+
+乱数を考えると非効率な状況もある（e.g. p= 1/2, 1）
+
+## step3
+詰まった点：
+- scaled_weights で n 倍するのを忘れる
+- 最後の while 処理を忘れる

0528.Random-Pick-with-Weight/step1_bisect.py

@@ -0,0 +1,19 @@
+import random
+import itertools
+import bisect
+
+class Solution:
+
+    def __init__(self, w: list[int]):
+        self.cumsum = list(itertools.accumulate(w))
+
+
+    def pickIndex(self) -> int:
+        sampled = random.uniform(0, self.cumsum[-1])
+        return bisect.bisect_left(self.cumsum, sampled)
+
+
+
+# Your Solution object will be instantiated and called as such:
+# obj = Solution(w)
+# param_1 = obj.pickIndex()

0528.Random-Pick-with-Weight/step1_naive.py

@@ -0,0 +1,20 @@
+import random
+import itertools
+
+class Solution:
+
+    def __init__(self, w: list[int]):
+        self.cumsum = list(itertools.accumulate(w))
+
+
+    def pickIndex(self) -> int:
+        sampled = random.uniform(0, self.cumsum[-1])
+        for i in range(len(self.cumsum)):
+            if sampled <= self.cumsum[i]:
+                return i
+
+
+
+# Your Solution object will be instantiated and called as such:
+# obj = Solution(w)
+# param_1 = obj.pickIndex()

0528.Random-Pick-with-Weight/step2_alias.py

@@ -0,0 +1,39 @@
+import random
+
+class Solution:
+
+    def __init__(self, weights: list[int]):
+        n = len(weights)
+        total = sum(weights)
+        scaled_prob = [w * n / total for w in weights]
+
+        self.prob = [0.0] * n
+        self.alias = [0] * n
+        self.n = n
+
+        small, large = [], []
+        for i, p in enumerate(scaled_prob):
+            if p < 1.0:
+                small.append(i)
+            else:
+                large.append(i)
+
+        while small and large:
+            s = small.pop()
+            l = large.pop()
+            self.prob[s] = scaled_prob[s]
+            self.alias[s] = l
+            scaled_prob[l] -= 1.0 - scaled_prob[s]
+            if scaled_prob[l] < 1.0:
+                small.append(l)
+            else:
+                large.append(l)
+
+        while large:
+            self.prob[large.pop()] = 1.0
+        while small:
+            self.prob[small.pop()] = 1.0
+
+    def pickIndex(self) -> int:
+        i = random.randint(0, self.n - 1)
+        return i if random.random() < self.prob[i] else self.alias[i]

0528.Random-Pick-with-Weight/step2_bisect.py

@@ -0,0 +1,17 @@
+import random
+import itertools
+import bisect
+
+class Solution:
+
+    def __init__(self, weights: list[int]):
+        self.prefix_sums = list(itertools.accumulate(weights))
+
+    def pickIndex(self) -> int:
+        target_weight = random.randint(0, self.prefix_sums[-1])
+        return bisect.bisect_left(self.prefix_sums, target_weight)
+
+
+# Your Solution object will be instantiated and called as such:
+# obj = Solution(w)
+# param_1 = obj.pickIndex()

0528.Random-Pick-with-Weight/step3_alias.py

@@ -0,0 +1,42 @@
+import random
+
+class Solution:
+
+    def __init__(self, weights: list[int]):
+        self.n = len(weights)
+        total = sum(weights)
+        scaled_weights = [w / total * self.n for w in weights]
+
+        self.prob = [0.0] * self.n
+        self.alias = [0] * self.n
+
+        small = []
+        large = []
+        for i, p in enumerate(scaled_weights):
+            if p < 1.0:
+                small.append(i)
+            else:
+                large.append(i)
+
+        while small and large:
+            i_small = small.pop()
+            i_large =large.pop()
+            self.alias[i_small] = i_large
+            self.prob[i_small] =  scaled_weights[i_small]
+            scaled_weights[i_large] -= 1.0 - scaled_weights[i_small]
+            if scaled_weights[i_large] < 1.0:
+                small.append(i_large)
+            else:
+                large.append(i_large)
+
+        while large:
+            i = large.pop()
+            self.prob[i] = 1.0
+        while small:
+            i = small.pop()
+            self.prob[i] = 1.0
+
+
+    def pickIndex(self) -> int:
+        i = random.randint(0, self.n - 1)
+        return i if random.random() < self.prob[i] else self.alias[i]