red_packets.py

import random
import matplotlib.pyplot as plt
import numpy as np
import functools
import time


def trunc_gauss(mu, sigma, bottom, top):
    a = random.gauss(mu,sigma)
    while (bottom <= a <= top) == False:
        a = random.gauss(mu,sigma)
    return a


def draw_stats(run_counts = 10000):
    def handle_func(func):
        @functools.wraps(func)
        def handle_args(*args, **kwargs):
            result_lists = []
            t = time.clock()
            for i in xrange(0, run_counts):
                result = func(*args, **kwargs)
                result_lists.append(result)
            print func.__name__, time.clock() - t
            result_lists = [item for sublist in result_lists for item in sublist]
            num_bins = 100
            plt.hist(result_lists,bins=num_bins, normed=1, facecolor='green', alpha=0.5)
            plt.xlabel('moneys')
            plt.ylabel('frequency')
            plt.subplots_adjust(left=0.15)
            plt.title(r'$\mu=%s$, $\sigma=%s$'%(np.mean(result_lists), np.std(result_lists)))
            plt.show()
        return handle_args
    return handle_func
    

@draw_stats()
def generate_red_packets(all_moneys = 100, peoples_num = 10, money_up_boundry = 13,money_down_boundry = 7):
    tmp_money_up_boundry = money_up_boundry - money_down_boundry
    tmp_money_down_boundry = 0
    left_distributie_moneys = all_moneys - peoples_num * money_down_boundry
    assert(left_distributie_moneys > 0)
    mu = left_distributie_moneys / float(peoples_num)
    list_moneys = [money_down_boundry] * peoples_num
    for i in xrange(0,peoples_num - 1):
        max_left_moneys = (peoples_num - i - 1) * tmp_money_up_boundry 
        while True:
            get_moneys = trunc_gauss(mu,1,tmp_money_down_boundry,tmp_money_up_boundry)
            #get_moneys = random.uniform(tmp_money_down_boundry, tmp_money_up_boundry)
            if max_left_moneys >= (left_distributie_moneys - get_moneys) and left_distributie_moneys - get_moneys >0:
                break
            if tmp_money_down_boundry > left_distributie_moneys:
                tmp_money_down_boundry = left_distributie_moneys
        left_distributie_moneys = left_distributie_moneys - get_moneys
        list_moneys[i] += get_moneys

    list_moneys[peoples_num - 1] += left_distributie_moneys
    return  list_moneys


@draw_stats()
def generate_red_packets_2(all_moneys = 100, peoples_num = 10, money_up_boundry = 13,money_down_boundry = 7,unit=0.01,random_func = random.uniform):
    list_moneys = [money_down_boundry] * peoples_num
    count_n = (all_moneys - money_down_boundry * peoples_num) / unit
    for i in xrange(0,int(count_n)):
        while True:
            index = int(random_func(0, peoples_num))
            if list_moneys[index] < money_up_boundry:
                list_moneys[index] += unit
                break
    return list_moneys


#generate_red_packets(money_up_boundry = 12,money_down_boundry = 6)
generate_red_packets_2(money_up_boundry = 12,money_down_boundry = 6)