June challenge day 5: "Random Pick with Weight" correctness
max_coder__
78

Hi, so I've tried to create my own approach to the "Pick with weight" problem. If I understood correctly we have to sample indexes by given weights, so if I have for example [3,3,3] for the weights I'm picking indexes 0-2 uniformly. In my approach I'm using SumTree which has same kind of approach of sampling by weights from elements.
However my solution fails to pass the tests! Any help or ideas why this happens? Also I'm not understanding how the solution should be exactly same in the tests, if we're using random numbers???

from typing import List
import random
import numpy as np


class SumTreeNode():
    def __init__(self, l=None, r=None, w=0, i=0):
        self.r = r
        self.l = l
        self.w = w
        self.i = i

class Solution:

    def __init__(self, w: List[int]):
        self.w = w
        self.wlen = len(w)
        self.rootNode = self.sum_tree(0, len(self.w))
        self.sum = sum(self.w)

    def pickIndex(self) -> int:
        rn = random.randint(1, self.sum)
        index = self.pickNodeIndex(self.rootNode, rn)
        return index

    def pickNodeIndex(self, node, value):
        if node.l is None:
            return node.i
        else:
            if value <= node.l.w:
                return self.pickNodeIndex(node.l, value - node.r.w)
            else:
                return self.pickNodeIndex(node.r, value - node.l.w)

    def sum_tree(self, s, l):
        if l == 1:
            node = SumTreeNode(w=self.w[s], i=s)
            return node
        else:
            half = l//2
            node_l = self.sum_tree(s, half)
            node_r = self.sum_tree(s + half, l - half)
            node = SumTreeNode(w=node_l.w + node_l.w, i=None, l=node_l, r=node_r)
            return node



if __name__ == "__main__":
    arr = [1, 1, 1]
    nparr = np.array(arr)
    s = Solution(list(arr))
    probs = nparr/np.sum(nparr)
    print(probs)
    # [0.33333333 0.33333333 0.33333333]
    samples = []
    for i in range(10000):
        samples.append(s.pickIndex())

    probs_algo = []
    for i in range(len(nparr)):
        probs_algo.append(samples.count(i)/len(samples))
    print(probs_algo)
    # [0.3358, 0.3302, 0.334]
Comments (0)