Make Groups

make_groups.py

"""Make Groups

# Problem

https://www.reddit.com/r/Python/comments/7awlm1/

I am a teacher and every year we have overnight excursions where we have to
organize rooms for around 100 students however there are a number of
requirements when organise these rooms:

1. Boys and girls MUST be in seperate rooms.

2. Room sizes and amount of rooms will vary. e.g. This year we might have 4x '5
bed' rooms and 7x '4 bed' rooms, however, we may have restrictions on this so
all the boys rooms are together etc..

3. Each student allocates 3 students they would like to share with and we
promise to get them at least 1 of their friends.

I understand this is a lot of variables, but I wanted to know if this is
possible and something I can work towards.

## Solution

Since groups are split boys/girls, you really have two problems and if you can
solve one then you can solve the other.

"""

import csv

from itertools import islice
from pprint import pprint
from random import sample, seed, shuffle
from warnings import warn


def assign(group_to_size, name_to_friends):
    "Assign names to groups."
    names = sorted(name_to_friends.keys())
    groups = sorted(group_to_size.keys())
    name_to_group = {}
    group_to_names = {}

    # Randomly assign names to groups.

    shuffle(names)
    names_iterator = iter(names)

    for group in groups:
        group_names = list(islice(names_iterator, group_to_size[group]))
        group_to_names[group] = group_names
        name_to_group.update({name: group for name in group_names})

    # Swap group names with friends until everyone has a friend in their group.

    for attempt in range(int(1e6)):
        swapped = 0
        shuffle(names)

        for name in names:
            group = name_to_group[name]
            group_names = group_to_names[group]
            friends = name_to_friends[name]

            if any(friend in group_names for friend in friends):
                continue

            # Select first friend and rotate friend list.

            friend = friends.pop(0)
            friends.append(friend)

            # Lookup friend group names.

            friend_group = name_to_group[friend]
            friend_group_names = group_to_names[friend_group]

            # Move friend to group names.

            friend_group_names.remove(friend)
            group_names.append(friend)
            name_to_group[friend] = group

            # Move first group name to friend group.

            group_name = group_names.pop(0)
            friend_group_names.append(group_name)
            name_to_group[group_name] = friend_group

            swapped += 1

        if not swapped:
            print(' ' * 80, end='\r')
            print(f'Assigned names to groups in {attempt + 1} iterations.')
            break

        if attempt % 1000 == 0:
            message = f'Iterations: {attempt} Swaps: {"*" * swapped}'
            print(f'{message:<80}', end='\r')
    else:
        warn(f'could not assign names to groups in {attempt + 1} iterations')

    return group_to_names


def validate(group_to_names, group_to_size, name_to_friends):
    """Validate group-to-names mapping against sizes and friends mappings.

    Assert groups meet size requirements. Warn for names with no friend in
    their group.

    """
    for group, names in group_to_names.items():
        assert len(names) == group_to_size[group]

        for name in names:
            friends = name_to_friends[name]

            if not any(friend in names for friend in friends):
                warn(f'{name!r} has no friend {friends!r} in group {group!r}')


def make_fake_data(names=56, groups=8, size=7, subset=7):
    """Generate fake data.

    Names are two-digit numbers: ['00', '01', '02', ...]
    Groups are single letters: ['a', 'b', 'c', ...]
    Friends are sampled randomly from a subset of names.

    """
    assert names == groups * size
    assert names % subset == 0

    names = ['%02d' % num for num in range(names)]
    groups = list('abcdefghijklmnopqrstuvwxyz'[:groups])
    group_to_size = {group: size for group in groups}
    name_to_friends = {}

    for name in names:  # Setup friends mapping.
        offset = int(name) // subset * subset
        pals = set(names[offset:(offset + subset)]) - {name}
        name_to_friends[name] = sample(sorted(pals), 3)

    return group_to_size, name_to_friends


def create_data_file(filename='example.csv', **kwargs):
    """Create data file with `filename`.

    Calls `make_fake_data` with `kwargs`

    """
    group_to_size, name_to_friends = make_fake_data(**kwargs)
    rows = [['Kind', 'Name', 'Data']]

    for group, size in group_to_size.items():
        rows.append(['Group', group, str(size)])

    for name, friends in name_to_friends.items():
        rows.append(['Student', name] + friends)

    with open(filename, 'w') as writer:
        csv_writer = csv.writer(writer)
        csv_writer.writerows(rows)


def parse_data_file(filepath):
    """Parse data file at `filepath`.

    """
    group_to_size = {}
    name_to_friends = {}

    with open(filepath, 'r') as file_reader:
        csv_reader = csv.reader(file_reader)
        header = next(csv_reader)
        assert header[:3] == ['Kind', 'Name', 'Data']

        for kind, name, *data in csv_reader:
            if kind == 'Group':
                group_to_size[name] = int(data[0])
            else:
                assert kind == 'Student'
                name_to_friends[name] = data

    return group_to_size, name_to_friends


def main(filepath):
    "Parse `filepath`, assign names to groups, validate and print."
    if filepath is None:
        filepath = input('CSV File Path: ')
        group_to_size, name_to_friends = parse_data_file(filepath)
    elif filepath == '-':
        group_to_size, name_to_friends = make_fake_data()
    else:
        group_to_size, name_to_friends = parse_data_file(filepath)

    group_to_names = assign(group_to_size, name_to_friends)
    validate(group_to_names, group_to_size, name_to_friends)
    pprint(group_to_names)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--random-seed', default=1)
    parser.add_argument('filepath', nargs='?', default=None)
    args = parser.parse_args()

    seed(args.random_seed)
    main(args.filepath)