solve.py

from __future__ import print_function
import sys
from ortools.constraint_solver import pywrapcp


def main():
    # Creates the solver.
    solver = pywrapcp.Solver("schedule_shifts")

    num_nurses = 4
    num_shifts = 4     # Nurse assigned to shift 0 means not working that day.
    num_days = 7

    # [START]
    # Create shift variables.
    shifts = {}

    for j in range(num_nurses):
        for i in range(num_days):
            shifts[(j, i)] = solver.IntVar(0, num_shifts - 1,
                                           "shifts(%i,%i)" % (j, i))

    shifts_flat = [shifts[(j, i)] for j in range(num_nurses)
                   for i in range(num_days)]

    # Create nurse variables.
    nurses = {}

    for j in range(num_shifts):
        for i in range(num_days):
            nurses[(j, i)] = solver.IntVar(0, num_nurses - 1,
                                           "shift%d day%d" % (j, i))

    nurses_flat = [nurses[(j, i)] for j in range(num_shifts)
                   for i in range(num_days)]

    # Set relationships between shifts and nurses.
    for day in range(num_days):
        nurses_for_day = [nurses[(j, day)] for j in range(num_shifts)]

        for j in range(num_nurses):
            s = shifts[(j, day)]
            solver.Add(s.IndexOf(nurses_for_day) == j)

    # Make assignments different on each day
    for i in range(num_days):
        solver.Add(solver.AllDifferent([shifts[(j, i)]
                                        for j in range(num_nurses)]))
        solver.Add(solver.AllDifferent([nurses[(j, i)]
                                        for j in range(num_shifts)]))

    # Each nurse works 5 or 6 days in a week.
    for j in range(num_nurses):
        solver.Add(solver.Sum([shifts[(j, i)] > 0
                               for i in range(num_days)]) >= 5)
        solver.Add(solver.Sum([shifts[(j, i)] > 0
                               for i in range(num_days)]) <= 6)

    # Create works_shift variables. works_shift[(i, j)] is True if nurse
    # i works shift j at least once during the week.
    works_shift = {}

    for i in range(num_nurses):
        for j in range(num_shifts):
            works_shift[(i, j)] = solver.BoolVar('shift%d nurse%d' % (i, j))

    for i in range(num_nurses):
        for j in range(num_shifts):
            solver.Add(works_shift[(i, j)] == solver.Max([shifts[(i, k)] == j for k in range(num_days)]))

    # For each shift (other than 0), at most 2 nurses are assigned to that
    # shift during the week.
    for j in range(1, num_shifts):
        solver.Add(solver.Sum([works_shift[(i, j)]
                              for i in range(num_nurses)]) <= 2)

    # If 's' nurse works shifts 2 or 3 on, he must also work that shift the
    # previous day or the following day.
    solver.Add(solver.Max(nurses[(2, 0)] == nurses[(2, 1)], nurses[(2, 1)] == nurses[(2, 2)]) == 1)
    solver.Add(solver.Max(nurses[(2, 1)] == nurses[(2, 2)], nurses[(2, 2)] == nurses[(2, 3)]) == 1)
    solver.Add(solver.Max(nurses[(2, 2)] == nurses[(2, 3)], nurses[(2, 3)] == nurses[(2, 4)]) == 1)
    solver.Add(solver.Max(nurses[(2, 3)] == nurses[(2, 4)], nurses[(2, 4)] == nurses[(2, 5)]) == 1)
    solver.Add(solver.Max(nurses[(2, 4)] == nurses[(2, 5)], nurses[(2, 5)] == nurses[(2, 6)]) == 1)
    solver.Add(solver.Max(nurses[(2, 5)] == nurses[(2, 6)], nurses[(2, 6)] == nurses[(2, 0)]) == 1)
    solver.Add(solver.Max(nurses[(2, 6)] == nurses[(2, 0)], nurses[(2, 0)] == nurses[(2, 1)]) == 1)

    solver.Add(solver.Max(nurses[(3, 0)] == nurses[(3, 1)], nurses[(3, 1)] == nurses[(3, 2)]) == 1)
    solver.Add(solver.Max(nurses[(3, 1)] == nurses[(3, 2)], nurses[(3, 2)] == nurses[(3, 3)]) == 1)
    solver.Add(solver.Max(nurses[(3, 2)] == nurses[(3, 3)], nurses[(3, 3)] == nurses[(3, 4)]) == 1)
    solver.Add(solver.Max(nurses[(3, 3)] == nurses[(3, 4)], nurses[(3, 4)] == nurses[(3, 5)]) == 1)
    solver.Add(solver.Max(nurses[(3, 4)] == nurses[(3, 5)], nurses[(3, 5)] == nurses[(3, 6)]) == 1)
    solver.Add(solver.Max(nurses[(3, 5)] == nurses[(3, 6)], nurses[(3, 6)] == nurses[(3, 0)]) == 1)
    solver.Add(solver.Max(nurses[(3, 6)] == nurses[(3, 0)], nurses[(3, 0)] == nurses[(3, 1)]) == 1)

    # Create the decision builder.
    db = solver.Phase(nurses_flat, solver.CHOOSE_FIRST_UNBOUND,
                      solver.ASSIGN_MIN_VALUE)

    # Create the solution collector.
    solution = solver.Assignment()
    solution.Add(nurses_flat)
    collector = solver.AllSolutionCollector(solution)

    solver.Solve(db, [collector])
    print("Solutions found:", collector.SolutionCount())
    print("Time:", solver.WallTime(), "ms")
    print()

    # Display a few solutions picked at random.
    def print_solution(sol):
        for i in range(num_days):
            print("Day", i)
            for j in range(num_nurses):
                print("Nurse", j, "assigned to shift",
                      collector.Value(sol, shifts[(j, i)]))
            print()

    def print_schedule(sol):
        days = ['Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat']
        nurse_dict = {0: 'A', 1: 'B', 2: 'C', 3: 'D'}

        row_format = '{:>7}' + '{:>5}' * num_days
        print(row_format.format('', *days))

        for j in range(1, num_shifts):
            row = []
            for i in range(num_days):
                val = collector.Value(sol, nurses[(j, i)])
                row.append(nurse_dict[val])
            print(row_format.format('Shift %d' % j, *row))
        print()

    a_few_solutions = [859, 2034, 5091, 7003]

    for sol in a_few_solutions:
        print("Solution number", sol, '\n')
        print_schedule(sol)

if __name__ == "__main__":
    main()