Simple students repartition solver

eskola.py

#!/usr/bin/python3
# Algo to find solutions to a simple classrooms organization problem
# Usage:
#   $ python3 eskola.py > out.txt

import sys
from enum import IntEnum
from operator import attrgetter
from random import choice, randrange, shuffle


class Level(IntEnum):
    CP = 1
    CE1 = 2
    CE2 = 3
    CM1 = 4
    CM2 = 5


STUDENTS = {
    Level.CP: 38,
    Level.CE1: 43,
    Level.CE2: 49,
    Level.CM1: 34,
    Level.CM2: 45,
}

MAX_CLASSROOMS = 8
MAX_SOLUTIONS = 10
MAX_STUDENTS_PER_CLASS = 28
MIN_STUDENTS_PER_CLASS = 20
MIN_STUDENTS_PER_LEVEL = 8


class RoomError(ValueError):
    pass


class Classroom():
    def __init__(self,  levels: list):
        self.repartition = {}
        for l in levels:
            self.repartition[l] = 0

    @property
    def levels(self):
        return tuple(sorted(self.repartition.keys()))

    @property
    def size(self):
        return sum(self.repartition.values())

    def add(self, l: Level):
        if l in self.repartition.keys() and self.size < MAX_STUDENTS_PER_CLASS:
            self.repartition[l] += 1
            return True
        return False

    def __str__(self):
        return ", ".join(map(lambda t: "{c:>2} {l.name:<3}".format(l=t[0], c=t[1]), self.repartition.items()))


def random_levels():
    lvl = choice([l.value for l in Level])
    if lvl == 5 or randrange(2):
        return [Level(lvl)]
    return [Level(lvl), Level(lvl+1)]


def find_classroom(l: Level, classrooms: list):
    for c in classrooms:
        if c.add(l):
            return
    raise RoomError()


def is_valid(classrooms: list):
    checkers = dict(STUDENTS)
    for c in classrooms:
        if c.size < MIN_STUDENTS_PER_CLASS:
            return False
        for l, c in c.repartition.items():
            if c < MIN_STUDENTS_PER_LEVEL:
                return False
            checkers[l] -= c
    if sum(checkers.values()) != 0:
        raise ValueError()
    return True


if __name__ == "__main__":
    # Keep solution to avoid doublons
    solutions = []
    while len(solutions) < MAX_SOLUTIONS:
        try:
            # Initialize all classrooms
            classrooms = []
            for _ in range(MAX_CLASSROOMS):
                classrooms.append(Classroom(random_levels()))

            # Initialize students list and shuffle it
            students = []
            for level, count in STUDENTS.items():
                students += [level] * count
            shuffle(students)
            for s in students:
                find_classroom(s, classrooms)

            # Sort solution to avoid doublons
            classrooms = sorted(classrooms, key=attrgetter("levels"))
            solution = tuple(map(attrgetter("levels"), classrooms))

            # Check solution not already found and valid
            if solution not in solutions and is_valid(classrooms):
                # Keep solution
                solutions.append(solution)
                print(len(solutions), end="", flush=True, file=sys.stderr)
                # Write solution to stdout
                print("Répartition {i}".format(i=len(solutions)))
                for c in classrooms:
                    print("\t", c)
                print(flush=True)

        except RoomError:
            pass
        finally:
            print(".", end="", flush=True, file=sys.stderr)