schedule_generator.py

"""
get permutations of schedules and assign students to them
"""

import csv
from itertools import permutations
from random import choice

MAX_CLASS_SIZE = 27

CLASSES = [
    'Geometry - Mr. Perkins (C222)',
    'W. History - Ms. Kaufman (C220)',
    'English 2 - Ms. Martinez (C218)',
    'Maritime - Mr. Moriarty (C224)',
    'Chemistry - Mr. Kowshik (D211)'
]


def flatten(l):
    """
    from: https://stackoverflow.com/questions/952914/making-a-flat-list-out-of-list-of-lists-in-python
    this exists because i'm lazy
    """
    flat_list = []
    for sublist in l:
        if isinstance(sublist, list) or isinstance(sublist, tuple):
            for item in sublist:
                flat_list.append(item)
        else:
            flat_list.append(sublist)

    return flat_list


def main():
    schedules = list(permutations(CLASSES))

    # fill our class counts
    counts = []
    for i in xrange(5):
        counts.append({c: 0 for c in CLASSES})

    students = []
    with open('MA10.csv', 'rb') as FILE:
        reader = csv.reader(FILE)
        reader.next()  # skip the header
        for row in reader:
            students.append({
                'name': row[0],
                'id': row[1],
                'mentor': row[2],
                'special': row[8] == '1'
            })

    # special kids to the front of the bus!
    students.sort(key=lambda k: k['special'], reverse=True)

    def will_overflow(schedule):
        for i, class_ in enumerate(schedule):
            # don't want imbalanced class sizes
            if capped(i) and counts[i][class_] + 1 == MAX_CLASS_SIZE:
                return True

            if counts[i][class_] + 1 > MAX_CLASS_SIZE:
                return True

        return False

    def capped(i):
        cap_num = len(students) % len(CLASSES)
        return sorted(counts[i].values())[-cap_num:] == [MAX_CLASS_SIZE for i in xrange(cap_num)]

    def increment_class(schedule):
        for i, class_ in enumerate(schedule):
            counts[i][class_] += 1

    for student in students:
        while not student.get('schedule'):
            # throws IndexError if it's hit a corner
            schedule = choice(schedules)

            if student['special'] and schedule[0] != CLASSES[0]:
                continue

            if will_overflow(schedule):
                schedules.remove(schedule)
                continue

            student['schedule'] = schedule
            increment_class(schedule)

    with open('MA10_output.csv', "wb+") as csv_file:
        writer = csv.writer(csv_file)
        writer.writerow(['ID', 'Name', 'Mentor', 'Period 1', 'Period 2',
                         'Period 3', 'Period 4', 'Period 5', 'Special?'])
        for student in students:
            writer.writerow(flatten([
                student['id'],
                student['name'],
                student['mentor'],
                student['schedule'],  # was a tuple, need to make a list
                student['special']
            ]))


count = 0
ATTEMPTS = 20
print 'finding a schedule!\n'

while count < ATTEMPTS:
    try:
        main()
        print 'found a schedule after', count, 'unsuccessful attempts'
        break
    except IndexError:
        count += 1

if count == 20:
    print 'could not find a valid schedule after', ATTEMPTS, 'tries. Adjust parameters and try again?',