A weighted version of random.choice

后端 未结 25 1976
闹比i
闹比i 2020-11-21 06:29

I needed to write a weighted version of random.choice (each element in the list has a different probability for being selected). This is what I came up with:



        
25条回答
  •  Happy的楠姐
    2020-11-21 06:44

    Provide random.choice() with a pre-weighted list:

    Solution & Test:

    import random
    
    options = ['a', 'b', 'c', 'd']
    weights = [1, 2, 5, 2]
    
    weighted_options = [[opt]*wgt for opt, wgt in zip(options, weights)]
    weighted_options = [opt for sublist in weighted_options for opt in sublist]
    print(weighted_options)
    
    # test
    
    counts = {c: 0 for c in options}
    for x in range(10000):
        counts[random.choice(weighted_options)] += 1
    
    for opt, wgt in zip(options, weights):
        wgt_r = counts[opt] / 10000 * sum(weights)
        print(opt, counts[opt], wgt, wgt_r)
    

    Output:

    ['a', 'b', 'b', 'c', 'c', 'c', 'c', 'c', 'd', 'd']
    a 1025 1 1.025
    b 1948 2 1.948
    c 5019 5 5.019
    d 2008 2 2.008
    

提交回复
热议问题