Randomized block assignment in Python

blocks.py

from __future__ import division
import random
import pprint

def assign_blocks(total_n, block_sizes):
    block_size_list = []
    current_n = 0
    while current_n < total_n:
        max_size = total_n - current_n
        possible_sizes = [s for s in block_sizes if s <= max_size]
        current_size = random.choice(possible_sizes)
        block_size_list.append(current_size)
        current_n += current_size
    print("Block sizes:")
    pprint.pprint(block_size_list)

    final_assignments = []
    for size in block_size_list:
        current_treatments = create_block(size)
        final_assignments.extend(current_treatments)
    return final_assignments

def weighted_choice(choices):
   total = sum(w for c, w in choices)
   r = random.uniform(0, total)
   upto = 0
   for c, w in choices:
      if upto + w > r:
         return c
      upto += w
   assert False, "Shouldn't get here"

def create_block(size):
    N_A = size / 2
    N_B = size / 2
    n_a = 0
    n_b = 0
    treatments = []
    while len(treatments) < size:
        p_a = (N_A - n_a) / (N_A + N_B - n_a - n_b)
        p_b = (N_B - n_b) / (N_A + N_B - n_a - n_b)
        current_treatment = weighted_choice([('A', p_a), ('B', p_b)])
        treatments.append(current_treatment)
        if current_treatment == 'A':
            n_a += 1
        if current_treatment == 'B':
            n_b += 1
    return treatments