Script for fairly distributing event tickets to people who joined in group order

For the last several years I’ve coordinated the group ticket order for my kids’ school for an annual dance festival that the school participates in. We don’t get to pick our seats for the group order; we tell the festival how many tickets we’re buying, and we get back tickets they’ve chosen. I’ve always tried to assign seats fairly, meaning family having about the same odds of getting good seats. This was easy to do by hand in previous years, but this year, that was much harder to do, because the order was much larger, a lot more people requested aisle seats, and the seats we got were in two qualitatively different groups (“close up but off to the side” vs. “centered but farther way”). So, I asked everyone if they had a preference of close up vs. centered, and then I wrote a script to do the seat assignments automatically. I present it here on the off chance that it will be useful to someone else. Comment below if you find it useful! If anybody actually comments then I will post updated versions of the script and comment when I do so you can subscribe to the RSS feed to find out about updates.

#!/usr/bin/env python3

# This script fairly, randomly distributes event tickets to a group of people
# who participated in a group ticket order and who ordered different quantities
# of tickets.
#
# The script reads a YAML file which defines:
#
# * rows of contiguous seats;
# * attributes associated with the rows and/or individual seats within them;
# * how many tickets each person ordered, along with attribute preferences; and
# * weights which indicate how significant each preference type is.
#
# When assigning tickets to individuals, the script works its way down from the
# people who ordered the most tickets to the people who ordered the least. This
# is necessary because otherwise the smaller orders break up blocks of
# contiguous seats and there's no place to put the larger orders when the
# script gets to them.
#
# The script assumes that seat numbers consist of a single or doubled letter
# followed by a number. If that's not the case, you may need to modify the
# 'is_adjacent', 'row', 'adjacent_rows', and/or 'seat_numbers' functions.
#
# The script runs continuously until it reaches a score of 0, i.e., it finds a
# seating arrangement that satisfies everybody's preferences. If you hit
# Ctrl-C it prints the current best set of assignments with their warnings. If
# you hit Ctrl-C twice in a row when a new best score was not generated in the
# interim, the script exits.
#
# The script tries to split orders of eight or more seats across two adjacent
# rows. If you don't want it to do that, make 'min_split' below very large.
# Make it smaller if you would like to split smaller orders as well.
#
# Example YAML input format:
#
# rows:
# - attributes: [close]
#   seat_attributes:
#     H24: [aisle]
#   seats: [H24, H26, H28, H30]
# - attributes: [centered]
#   seats: [P16, P18, P20, P22]
# people:
#   example@example.com:
#     num_tickets: 2
#     preferences: [centered]
#     seat_preferences: [aisle]
#     assign: [P20, P22]
# weights:
#   aisle: 100
#   centered: 1
#   close: 1
#
# 'attributes', 'seat_attributes', 'preferences', 'seat_preferences', and
# 'assign' are optional. 'assign' is used to override the random logic and
# assign specific seats to someone. Contents of 'attributes' should match with
# contents of 'preferences' and similarly 'seat_attributes' should match with
# 'seat_preferences'. 'weights' are optional and tell the script how important
# to consider each preference when calculating the score of a particular set
# of seating assignments.
#
# Copyright 2019 Jonathan Kamens . Feel free to email me
# comments or questions.
#
# This script is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# This script is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
# details.
#
# A copy of the GNU General Public License can be found at
# https://www.gnu.org/licenses/.

import argparse
from collections import defaultdict
import copy
import pprint
import random
import re
import signal
import sys
import yaml

min_split = 8
best = {'score': 99999999}


def parse_args():
    parser = argparse.ArgumentParser(description='Assign tickets to people '
                                     'based on preferences')
    parser.add_argument('data_file', metavar='DATA_FILE',
                        help='YAML file containing ticket and people data')
    args = parser.parse_args()
    return args


def read_data(args):
    with open(args.data_file, 'r') as f:
        data = yaml.load(f)
    row_attributes = set()
    seat_attributes = set()
    row_preferences = set()
    seat_preferences = set()

    assert 'rows' in data
    assert 'people' in data
    data['weights'] = data.get('weights', {})
    assert not any(k for k in data if k not in ('rows', 'people', 'weights'))

    for i in range(len(data['rows'])):
        row = data['rows'][i]
        row['index'] = i
        row['attributes'] = row.get('attributes', {})
        row['seat_attributes'] = row.get('seat_attributes', {})
        if 'seats' not in row:
            raise Exception('No seats in row {}'.format(row))
        for key, value in row.items():
            if key == 'attributes':
                row_attributes.update(value)
            elif key == 'seat_attributes':
                for seat, attributes in value.items():
                    assert seat in row['seats']
                    seat_attributes.update(attributes)
            elif key in ('seats', 'index'):
                pass
            else:
                raise Exception('Unrecognized row field: {}'.format(key))
        row['unassigned'] = [list(row['seats'])]

    for name, params in data['people'].items():
        params['preferences'] = params.get('preferences', {})
        params['seat_preferences'] = params.get('seat_preferences', {})
        if 'num_tickets' not in params:
            raise Exception('No num_tickets for {}'.format(name))
        for key, value in params.items():
            if key == 'num_tickets':
                pass
            elif key == 'preferences':
                row_preferences.update(value)
            elif key == 'seat_preferences':
                seat_preferences.update(value)
            elif key == 'assign':
                if len(value) != params['num_tickets']:
                    raise Exception('Wrong number of seats assigned for {}'.
                                    format(name))
                for s in value:
                    if not any(r for r in data['rows'] if s in r['seats']):
                        raise Exception('Invalid seat {} for {}'.format(
                            s, name))
            else:
                raise Exception('Unrecognized person field {} for {}'.format(
                    key, name))
        # Doing this here rather than above because I don't want to modify
        # params while iterating through it
        if 'assign' in params:
            assign(data, name, params['assign'])

    extra_row_attributes = row_attributes - row_preferences
    extra_row_preferences = row_preferences - row_attributes
    extra_seat_attributes = seat_attributes - seat_preferences
    extra_seat_preferences = seat_preferences - seat_attributes

    if extra_row_preferences:
        raise Exception('No rows with attribute(s) {}'.format(
            extra_row_preferences))
    if extra_seat_preferences:
        raise Exception('No seats with attribute(s) {}'.format(
            extra_seat_preferences))
    if extra_row_attributes:
        sys.stderr.write('Warning: No people with row attribute(s) {}\n'.
                         format(extra_row_attributes))
    if extra_seat_attributes:
        sys.stderr.write('Warning: no people with seat attribute(s) {}\n'.
                         format(extra_seat_attributes))

    data['names'] = list(data['people'].keys())

    return data


def sigint_handler(sig, frame):
    print_assignments()
    signal.signal(signal.SIGINT, signal.SIG_DFL)


def main():
    global max_attempts, current_attempts
    args = parse_args()
    data = read_data(args)
    max_attempts = len(data['names'])

    data['names'].sort(key=lambda n: data['people'][n]['num_tickets'],
                       reverse=True)
    while True:
        current_attempts = 0
        shuffle_rows(data)
        if recurse(data):
            break


def shuffle_rows(data):
    random.shuffle(data['rows'])
    for i in range(len(data['rows'])):
        data['rows'][i]['index'] = i


def recurse(data):
    global current_attempts

    names = [n for n in data['names']
             if 'assigned' not in data['people'][n]]
    if not names:
        return save_assignments(data)

    current_attempts += 1
    for name in names:
        for seats, row in valid_assignments(data, name):
            n_data = copy.deepcopy(data)
            assign(n_data, name, seats, row)
            if recurse(n_data):
                return True
            if current_attempts >= max_attempts:
                return False
    return False


def find_ticket_blocks(data, num_tickets, adjacent_to=None):
    for row in data['rows']:
        for unassigned in row['unassigned']:
            for i in range(0, len(unassigned) - num_tickets + 1):
                seats = unassigned[i:i+num_tickets]
                if not adjacent_to or is_adjacent(seats, adjacent_to):
                    yield seats, row


def ticket_block_score(data, name, seats, row):
    # Lower score is better
    score = 2**16
    for p in data['people'][name]['preferences']:
        if p in row['attributes']:
            score -= data['weights'].get(p, 1)
    for p in data['people'][name]['seat_preferences']:
        if any(p in row['seat_attributes'].get(s, []) for s in seats):
            score -= data['weights'].get(p, 1)
    return score


def valid_assignments(data, name):
    num_tickets = data['people'][name]['num_tickets']
    if num_tickets >= min_split:
        # Prefer to split large groups across two rows, to give them a
        # better chance of getting decent seats and make the whole thing
        # run faster.
        num_first = int(num_tickets / 2)
        num_second = num_tickets - num_first
        first = sorted(
            find_ticket_blocks(data, num_first),
            key=lambda t: ticket_block_score(data, name, t[0], t[1]))
        for first_seats, row in first:
            second = find_ticket_blocks(data, num_second, first_seats)
            for second_seats, row in second:
                yield first_seats + second_seats, None
    blocks = sorted(
        find_ticket_blocks(data, num_tickets),
        key=lambda t: ticket_block_score(data, name, t[0], t[1]))
    for seats, row in blocks:
        yield seats, row['index']


def is_adjacent(block1, block2):
    row1 = row(block1[0])
    row2 = row(block2[0])
    if row1 == row2:
        return False
    if not adjacent_rows(row1, row2):
        return False
    seats1 = set(seat_numbers(block1))
    seats2 = set(seat_numbers(block2))
    overlap = seats1 & seats2
    return len(seats1 - overlap) < 2 and len(seats2 - overlap) < 2


def row(seat):
    match = re.match(r'[A-Z]+', seat)
    return match.group(0)


def adjacent_rows(row1, row2):
    if len(row1) != len(row2):
        return False
    return abs(ord(row1[0]) - ord(row2[0])) == 1


def seat_numbers(seats):
    for seat in seats:
        match = re.search(r'(\d+)$', seat)
        yield match.group(0)


def assign(data, name, seats, row_number=None):
    if row_number is None:
        found_rows = []
        for seat in seats:
            try:
                existing_row = next(r for r in found_rows
                                    if seat in r['row']['seats'])
                existing_row['seats'].append(seat)
            except StopIteration:
                row = next(r for r in data['rows'] if seat in r['seats'])
                found_rows.append({'seats': [seat], 'row': row})
        for found_row in found_rows:
            assign(data, name, found_row['seats'], found_row['row']['index'])
        return

    row = data['rows'][row_number]
    person = data['people'][name]
    if 'assigned' not in person:
        person['assigned'] = seats
    else:
        person['assigned'].extend(seats)

    for i in range(len(row['unassigned']) - 1, -1, -1):
        before = []
        during = []
        after = []
        for row_seat in row['unassigned'][i]:
            if row_seat in seats:
                during.append(row_seat)
            elif during:
                after.append(row_seat)
            else:
                before.append(row_seat)
        if before and after:
            row['unassigned'][i:i+1] = [before, after]
        elif before:
            row['unassigned'][i] = before
        elif after:
            row['unassigned'][i] = after
        else:
            row['unassigned'][i:i+1] = []


def save_assignments(data):
    global best

    score = 0
    warnings = ''
    satisfied = True
    for name, params in data['people'].items():
        rows = set(r['index'] for r in data['rows']
                   for s in r['seats']
                   if s in params['assigned'])
        row_attributes = set(a for r in rows
                             for a in data['rows'][r]['attributes'])

        seat_attributes = set(a for r in rows
                              for s in params['assigned']
                              for a in data['rows'][r]['seat_attributes'].get(
                                      s, []))
        missing_preferences = set(params['preferences']) - row_attributes
        missing_seat_preferences = set(params['seat_preferences']) - \
            seat_attributes
        missing = missing_preferences | missing_seat_preferences
        for pref in missing:
            satisfied = False
            score += data['weights'].get(pref, 0)
            warnings += ('Warning: failed to satisfy preference(s) {} '
                         'for {}\n'.format(missing, name))

    if score < best['score']:
        del data['names']
        for params in data['people'].values():
            if not params['preferences']:
                del params['preferences']
            if not params['seat_preferences']:
                del params['seat_preferences']
        for row in data['rows']:
            del row['index']
            del row['unassigned']
            if not row['attributes']:
                del row['attributes']
            if not row['seat_attributes']:
                del row['seat_attributes']
        best['score'] = score
        best['warnings'] = warnings
        best['data'] = data
        print('New best score is {}'.format(score))
        signal.signal(signal.SIGINT, sigint_handler)

    if satisfied:
        print_assignments()
    return satisfied


def print_assignments():
    print('')
    print(best['warnings'])
    data = best['data']
    pprint.pprint(data, compact=True)
    print('')
    for name in sorted(data['people'].keys()):
        for seat in data['people'][name]['assigned']:
            print('{} {}'.format(name, seat))


if __name__ == '__main__':
    main()