Apportionment algorithm for the US House of Representatives

apportion.py

# Uses method of least proportions to apportion House members among states.
# http://en.wikipedia.org/wiki/United_States_congressional_apportionment#The_Method_of_Equal_Proportions
# use like this: 
#   python apportion.py --built-in 2010
#   python apportion.py --file file.csv --key <column name>
#   python apportion.py --file file.csv --key <column name> --reps 444
#   python apportion.py --file file.csv --key <column name> --distsize 500000
from __future__ import division, print_function
import sys
import csv
import argparse


BUILTINS = {
    # 2010 Apportionment populations
    '2010': {
        "Wyoming": 568300,
        "Vermont": 630337,
        "North Dakota": 675905,
        "Alaska": 721523,
        "South Dakota": 819761,
        "Delaware": 900877,
        "Montana": 994416,
        "Rhode Island": 1055247,
        "New Hampshire": 1321445,
        "Maine": 1333074,
        "Hawaii": 1366862,
        "Idaho": 1573499,
        "Nebraska": 1831825,
        "West Virginia": 1859815,
        "New Mexico": 2067273,
        "Nevada": 2709432,
        "Utah": 2770765,
        "Kansas": 2863813,
        "Arkansas": 2926229,
        "Mississippi": 2978240,
        "Iowa": 3053787,
        "Connecticut": 3581628,
        "Oklahoma": 3764882,
        "Oregon": 3848606,
        "Kentucky": 4350606,
        "Louisiana": 4553962,
        "South Carolina": 4645975,
        "Alabama":  4802982,
        "Colorado":  5044930,
        "Minnesota":  5314879,
        "Wisconsin":  5698230,
        "Maryland":  5789929,
        "Missouri":  6011478,
        "Tennessee":  6375431,
        "Arizona":  6412700,
        "Indiana":  6501582,
        "Massachusetts":  6559644,
        "Washington":  6753369,
        "Virginia":  8037736,
        "New Jersey":  8807501,
        "North Carolina":  9565781,
        "Georgia":  9727566,
        "Michigan":  9911626,
        "Ohio":  11568495,
        "Pennsylvania":  12734905,
        "Illinois":  12864380,
        "Florida":  18900773,
        "New York":  19421055,
        "Texas":  25268418,
        "California": 37341989
    },

    # 2000 Apportionment populations
    '2000': {
        "Alabama": 4461130,
        "Alaska": 628933,
        "Arizona": 5140683,
        "Arkansas": 2679733,
        "California": 33930798,
        "Colorado": 4311882,
        "Connecticut": 3409535,
        "Delaware": 785068,
        "Florida": 16028890,
        "Georgia": 8206975,
        "Hawaii": 1216642,
        "Idaho": 1297274,
        "Illinois": 12439042,
        "Indiana": 6090782,
        "Iowa": 2931923,
        "Kansas": 2693824,
        "Kentucky": 4049431,
        "Louisiana": 4480271,
        "Maine": 1277731,
        "Maryland": 5307886,
        "Massachusetts": 6355568,
        "Michigan": 9955829,
        "Minnesota": 4925670,
        "Mississippi": 2852927,
        "Missouri": 5606260,
        "Montana": 905316,
        "Nebraska": 1715369,
        "Nevada": 2002032,
        "New Hampshire": 1238415,
        "New Jersey": 8424354,
        "New Mexico": 1823821,
        "New York": 19004973,
        "North Carolina": 8067673,
        "North Dakota": 643756,
        "Ohio": 11374540,
        "Oklahoma": 3458819,
        "Oregon": 3428543,
        "Pennsylvania": 12300670,
        "Rhode Island": 1049662,
        "South Carolina": 4025061,
        "South Dakota": 756874,
        "Tennessee": 5700037,
        "Texas": 20903994,
        "Utah": 2236714,
        "Vermont": 609890,
        "Virginia": 7100702,
        "Washington": 5908684,
        "West Virginia": 1813077,
        "Wisconsin": 5371210,
        "Wyoming": 495304
    }
}


class state(object):

    """A single state has a population and a current reps number"""

    def __init__(self, pop, name=None):
        self.name = name or ''
        self.pop = int(pop)
        self.reps = 1

    def priority(self):
        d = pow(self.reps * (self.reps + 1), 1 / 2)
        return self.pop / d

    def addseat(self):
        self.reps += 1

    def __repr__(self):
        return self.name + " (" + repr(self.pop) + "): " + repr(self.reps)


def allocated(states):
    return sum(s.reps for s in states)


def top_priority(states):
    '''Get top priority state from list of `State`s'''
    return max(states, key=lambda x: x.priority())


def read_csv(filename, key=None):
    with open(filename, 'r') as f:
        reader = csv.reader(f)
        header = next(reader)

        try:
            index = header.index(key)
        except ValueError:
            index = len(header)

        return {row[0]: float(row[index]) for row in reader}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--file', type=str)
    parser.add_argument('--key', type=str)
    parser.add_argument('--reps', type=int, default=435)
    parser.add_argument('--distsize', type=int, default=None)
    parser.add_argument('--built-in', type=str, default=None, choices=('2010', '2000'))

    args = parser.parse_args()

    # geographies is a dict with {name: population}.

    if args.built_in:
        geographies = BUILTINS[args.built_in]

    else:
        geographies = read_csv(args.file, args.key)

    total_pop = sum(geographies.values())

    if args.distsize:
        args.reps = round(total_pop / args.distsize)
    elif args.reps:
        args.sdistsize = round(total_pop / args.reps)
    else:
        raise ValueError("Need either a target district size or a number of reps!")

    states = {k: state(v, k) for k, v in geographies.items()}

    print("Total reps:", int(args.reps), file=sys.stderr)
    print("Avg district size:", int(args.distsize), file=sys.stderr)
    print("Created", len(states), "states", file=sys.stderr)

    while allocated(states.values()) < args.reps:
        top_priority(states.values()).addseat()

    # Sort alphabetically for output
    writer = csv.writer(sys.stdout, delimiter='\t')
    writer.writerow(['name', 'seats', 'pop'])
    for name in sorted(states.keys()):
        writer.writerow([name, states[name].reps, states[name].pop])


if __name__ == '__main__':
    main()