# Weighted versions of 'random.choice' 0.2.0
# Copyright 2003 Lja MultiMedia
# May be used under the GNU GPL

import time, random, bisect

def w_choice_a(seq):
   """original algorithm"""
   seq2 = []
   for (choice, weight) in seq:
      seq2 += [choice] * weight
   return random.choice(seq2)

def w_choice_b(seq):
   """exarkun's suggestion"""
   choices = []
   weights = []
   total = 0
   for (choice, weight) in seq:
      choices.append(choice)
      weights.append(weight)
      total += weight
   num = random.randrange(total)
   for i in range(len(weights)):
      num -= weights[i]
      if num <= 0:
         return choices[i]
   raise "This can't happen!"

def w_choice_c(seq):
   """levi's bisect idea"""
   choices = []
   weights = [0]
   for (choice, weight) in seq:
      choices.append(choice)
      weights.append(weights[-1] + weight)
   del weights[0]
   return choices[bisect.bisect(weights, random.randrange(weights[-1]))]

def _time(func, *args, **kwargs):
   start = time.clock()
   value = apply(func, args, kwargs)
   stop = time.clock()
   return (stop - start, value)

def _rep(num, func, *args, **kwargs):
   for i in range(num):
      apply(func, args, kwargs)

def _bench():
   seq0 = (('a',1),('b',2),('c',3))
   seq1 = (('a',186),('b',257),('c',193))
   seq2 = (('a',1876),('b',2547),('c',1983))
   seq3 = (('a',1876),('b',2547),('c',1983),('d',1476),('e',2573),('f',1933))
   
   for func in (w_choice_a, w_choice_b, w_choice_c):
      print '--', func.func_name, '--'
      print _time(_rep, 10000, func, seq0)
      print _time(_rep, 10000, func, seq1)
      print _time(_rep, 10000, func, seq2)
      print _time(_rep, 10000, func, seq3)

def _test():
   print 'Warning: No test, just this benchmark:'
   _bench()

w_choice = w_choice_b

if __name__ == '__main__':
   _test()