examples/notebook/nurses_sat.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2010-2017 Google\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "\n",
    "from __future__ import print_function\n",
    "import sys\n",
    "from ortools.sat.python import cp_model\n",
    "\n",
    "\n",
    "class NursesPartialSolutionPrinter(cp_model.CpSolverSolutionCallback):\n",
    "  \"\"\"Print intermediate solutions.\"\"\"\n",
    "\n",
    "  def __init__(self, shifts, num_nurses, num_days, num_shifts, sols):\n",
    "    cp_model.CpSolverSolutionCallback.__init__(self)\n",
    "    self.__shifts = shifts\n",
    "    self.__num_nurses = num_nurses\n",
    "    self.__num_days = num_days\n",
    "    self.__num_shifts = num_shifts\n",
    "    self.__solutions = set(sols)\n",
    "    self.__solution_count = 0\n",
    "\n",
    "  def OnSolutionCallback(self):\n",
    "    self.__solution_count += 1\n",
    "    if self.__solution_count in self.__solutions:\n",
    "      print('Solution #%i' % self.__solution_count)\n",
    "      for d in range(self.__num_days):\n",
    "        print('Day #%i' % d)\n",
    "        for n in range(self.__num_nurses):\n",
    "          for s in range(self.__num_shifts):\n",
    "            if self.Value(self.__shifts[(n, d, s)]):\n",
    "              print('  Nurse #%i is working shift #%i' % (n, s))\n",
    "      print()\n",
    "\n",
    "  def SolutionCount(self):\n",
    "    return self.__solution_count\n",
    "\n",
    "\n",
    "# Data.\n",
    "num_nurses = 4\n",
    "num_shifts = 4  # Nurse assigned to shift 0 means not working that day.\n",
    "num_days = 7\n",
    "all_nurses = range(num_nurses)\n",
    "all_shifts = range(num_shifts)\n",
    "all_working_shifts = range(1, num_shifts)\n",
    "all_days = range(num_days)\n",
    "\n",
    "# Creates the model.\n",
    "model = cp_model.CpModel()\n",
    "\n",
    "# Creates shift variables.\n",
    "# shifts[(n, d, s)]: nurse 'n' works shift 's' on day 'd'.\n",
    "shifts = {}\n",
    "for n in all_nurses:\n",
    "  for d in all_days:\n",
    "    for s in all_shifts:\n",
    "      shifts[(n, d, s)] = model.NewBoolVar('shift_n%id%is%i' % (n, d, s))\n",
    "\n",
    "# Makes assignments different on each day, that is each shift is assigned at\n",
    "# most one nurse. As we have the same number of nurses and shifts, then each\n",
    "# day, each shift is assigned to exactly one nurse.\n",
    "for d in all_days:\n",
    "  for s in all_shifts:\n",
    "    model.Add(sum(shifts[(n, d, s)] for n in all_nurses) == 1)\n",
    "\n",
    "# Nurses do 1 shift per day.\n",
    "for n in all_nurses:\n",
    "  for d in all_days:\n",
    "    model.Add(sum(shifts[(n, d, s)] for s in all_shifts) == 1)\n",
    "\n",
    "# Each nurse works 5 or 6 days in a week.\n",
    "# That is each nurse works shift 0 at most 2 times.\n",
    "for n in all_nurses:\n",
    "  model.AddSumConstraint([shifts[(n, d, 0)] for d in all_days], 1, 2)\n",
    "\n",
    "# works_shift[(n, s)] is 1 if nurse n works shift s at least one day in\n",
    "# the week.\n",
    "works_shift = {}\n",
    "for n in all_nurses:\n",
    "  for s in all_shifts:\n",
    "    works_shift[(n, s)] = model.NewBoolVar('works_shift_n%is%i' % (n, s))\n",
    "    model.AddMaxEquality(works_shift[(n, s)],\n",
    "                         [shifts[(n, d, s)] for d in all_days])\n",
    "\n",
    "# For each shift, at most 2 nurses are assigned to that shift during the week.\n",
    "for s in all_working_shifts:\n",
    "  model.Add(sum(works_shift[(n, s)] for n in all_nurses) <= 2)\n",
    "\n",
    "# If a nurse works shifts 2 or 3 on, she must also work that shift the\n",
    "# previous day or the following day.\n",
    "# This means that on a given day and shift, either she does not work that\n",
    "# shift on that day, or she works that shift on the day before, or the day\n",
    "# after.\n",
    "for n in all_nurses:\n",
    "  for s in [2, 3]:\n",
    "    for d in all_days:\n",
    "      yesterday = (d - 1) % num_days\n",
    "      tomorrow = (d + 1) % num_days\n",
    "      model.AddBoolOr([\n",
    "          shifts[(n, yesterday, s)], shifts[(n, d, s)].Not(),\n",
    "          shifts[(n, tomorrow, s)]\n",
    "      ])\n",
    "\n",
    "# Creates the solver and solve.\n",
    "solver = cp_model.CpSolver()\n",
    "# Display a few solutions picked at random.\n",
    "a_few_solutions = [859, 2034, 5091, 7003]\n",
    "solution_printer = NursesPartialSolutionPrinter(shifts, num_nurses, num_days,\n",
    "                                                num_shifts, a_few_solutions)\n",
    "status = solver.SearchForAllSolutions(model, solution_printer)\n",
    "\n",
    "# Statistics.\n",
    "print()\n",
    "print('Statistics')\n",
    "print('  - conflicts       : %i' % solver.NumConflicts())\n",
    "print('  - branches        : %i' % solver.NumBranches())\n",
    "print('  - wall time       : %f ms' % solver.WallTime())\n",
    "print('  - solutions found : %i' % solution_printer.SolutionCount())\n",
    "\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}