Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save bartaelterman/3954eff5e5249e52cd473eec724f544e to your computer and use it in GitHub Desktop.
Save bartaelterman/3954eff5e5249e52cd473eec724f544e to your computer and use it in GitHub Desktop.
Display the source blob
Display the rendered blob
Raw
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"from ortools.sat.python import cp_model\n",
"from random import choice"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"persons = ['Mam', 'Dad', 'Bart', 'Belle', 'Jonathan', 'Johanna', 'Alex', 'Alice', 'Robert', 'Rhonda']\n",
"\n",
"forbidden = [\n",
" ['Mam', 'Dad'],\n",
" ['Bart', 'Belle',],\n",
" ['Jonathan', 'Johanna'],\n",
" ['Alex', 'Alice'],\n",
" ['Robert', 'Rhonda']\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"model = cp_model.CpModel()"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"variables = {}\n",
"\n",
"for buyer in persons:\n",
" receiver_dict = {}\n",
" for receiver in persons:\n",
" receiver_dict[receiver] = model.NewIntVar(0, 1, f'{buyer} buys for {receiver}')\n",
" variables[buyer] = receiver_dict"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"for name, buyer in variables.items():\n",
" model.Add(sum(buyer.values()) == 1)\n",
"\n",
"for name in persons:\n",
" model.Add(sum([receivers[name] for buyer_name, receivers in variables.items()]) == 1)\n",
"\n",
"for name in persons:\n",
" model.Add(variables[name][name] == 0)\n",
"\n",
"for name_1, name_2 in forbidden:\n",
" model.Add(variables[name_1][name_2] == 0)\n",
" model.Add(variables[name_2][name_1] == 0)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"class AllSolutionsStore(cp_model.CpSolverSolutionCallback):\n",
" \"\"\"Count all solutions\"\"\"\n",
"\n",
" def __init__(self, variables):\n",
" cp_model.CpSolverSolutionCallback.__init__(self)\n",
" self.__variables = variables\n",
" self.__solution_count = 0\n",
"\n",
"\n",
" def on_solution_callback(self):\n",
" self.__solution_count += 1\n",
"\n",
" def solution_count(self):\n",
" return self.__solution_count"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"all_variables = []\n",
"for name, x in variables.items():\n",
" all_variables += list(x.values())"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"solver = cp_model.CpSolver()\n",
"solutionsstore = AllSolutionsStore(all_variables)\n",
"result = solver.SearchForAllSolutions(model, solutionsstore)"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Number of solutions found: 440192\n"
]
}
],
"source": [
"print('Number of solutions found: %i' % solutionsstore.solution_count())"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.6.9"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment