Skip to content

Instantly share code, notes, and snippets.

@hsharrison
Created July 12, 2019 09:02
Show Gist options
  • Save hsharrison/9a6734099af38be2f195efba44aafae9 to your computer and use it in GitHub Desktop.
Save hsharrison/9a6734099af38be2f195efba44aafae9 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": [
"class Kid:\n",
" def __init__(self, name, sport, camp=None):\n",
" self.name = name\n",
" self.sport = sport\n",
" self.camp = camp\n",
"\n",
"class Camp:\n",
" def __init__(self, name, sports):\n",
" self.name = name\n",
" self.sports = sports\n",
" self.open_spots = 3\n",
"\n",
"kids = [\n",
" Kid('Ben', 'Baseball'),\n",
" Kid('John', 'Baseball'),\n",
" Kid('Tom', 'Baseball'),\n",
" Kid('Chris', 'Soccer'),\n",
" Kid('Josh', 'Soccer'),\n",
" Kid('Alan', 'Basketball'),\n",
" Kid('Timmy', 'Basketball'),\n",
" Kid('Greg', 'Tennis'),\n",
"]\n",
"\n",
"camps = [\n",
" Camp('Cool Camp', ['Baseball', 'Soccer']),\n",
" Camp('Fun Camp', ['Baseball']),\n",
" Camp('Super Camp', ['Basketball']),\n",
" Camp('Sporty Camp', ['Basketball']),\n",
"]\n",
"\n",
"expected_answer = {\n",
" 'Ben': 'Cool Camp',\n",
" 'John': 'Fun Camp',\n",
" 'Tom': 'Fun Camp',\n",
" 'Chris': 'Cool Camp',\n",
" 'Josh': 'Cool Camp',\n",
" 'Alan': 'Super Camp',\n",
" 'Timmy': 'Sporty Camp',\n",
" 'Greg': None,\n",
"}\n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"import pandas as pd\n",
"import pulp"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"kids_by_name = {\n",
" kid.name: kid for kid in kids\n",
"}\n",
"camps_by_name = {\n",
" camp.name: camp for camp in camps\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"kids_names = list(kids_by_name)\n",
"camps_names = list(camps_by_name)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>Cool Camp</th>\n",
" <th>Fun Camp</th>\n",
" <th>Super Camp</th>\n",
" <th>Sporty Camp</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>Ben</th>\n",
" <td>Ben|Cool_Camp</td>\n",
" <td>Ben|Fun_Camp</td>\n",
" <td>Ben|Super_Camp</td>\n",
" <td>Ben|Sporty_Camp</td>\n",
" </tr>\n",
" <tr>\n",
" <th>John</th>\n",
" <td>John|Cool_Camp</td>\n",
" <td>John|Fun_Camp</td>\n",
" <td>John|Super_Camp</td>\n",
" <td>John|Sporty_Camp</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Tom</th>\n",
" <td>Tom|Cool_Camp</td>\n",
" <td>Tom|Fun_Camp</td>\n",
" <td>Tom|Super_Camp</td>\n",
" <td>Tom|Sporty_Camp</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Chris</th>\n",
" <td>Chris|Cool_Camp</td>\n",
" <td>Chris|Fun_Camp</td>\n",
" <td>Chris|Super_Camp</td>\n",
" <td>Chris|Sporty_Camp</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Josh</th>\n",
" <td>Josh|Cool_Camp</td>\n",
" <td>Josh|Fun_Camp</td>\n",
" <td>Josh|Super_Camp</td>\n",
" <td>Josh|Sporty_Camp</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Alan</th>\n",
" <td>Alan|Cool_Camp</td>\n",
" <td>Alan|Fun_Camp</td>\n",
" <td>Alan|Super_Camp</td>\n",
" <td>Alan|Sporty_Camp</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Timmy</th>\n",
" <td>Timmy|Cool_Camp</td>\n",
" <td>Timmy|Fun_Camp</td>\n",
" <td>Timmy|Super_Camp</td>\n",
" <td>Timmy|Sporty_Camp</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Greg</th>\n",
" <td>Greg|Cool_Camp</td>\n",
" <td>Greg|Fun_Camp</td>\n",
" <td>Greg|Super_Camp</td>\n",
" <td>Greg|Sporty_Camp</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" Cool Camp Fun Camp Super Camp Sporty Camp\n",
"Ben Ben|Cool_Camp Ben|Fun_Camp Ben|Super_Camp Ben|Sporty_Camp\n",
"John John|Cool_Camp John|Fun_Camp John|Super_Camp John|Sporty_Camp\n",
"Tom Tom|Cool_Camp Tom|Fun_Camp Tom|Super_Camp Tom|Sporty_Camp\n",
"Chris Chris|Cool_Camp Chris|Fun_Camp Chris|Super_Camp Chris|Sporty_Camp\n",
"Josh Josh|Cool_Camp Josh|Fun_Camp Josh|Super_Camp Josh|Sporty_Camp\n",
"Alan Alan|Cool_Camp Alan|Fun_Camp Alan|Super_Camp Alan|Sporty_Camp\n",
"Timmy Timmy|Cool_Camp Timmy|Fun_Camp Timmy|Super_Camp Timmy|Sporty_Camp\n",
"Greg Greg|Cool_Camp Greg|Fun_Camp Greg|Super_Camp Greg|Sporty_Camp"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"variables = pd.DataFrame(\n",
" [\n",
" [pulp.LpVariable(f'{kid}|{camp}', 0, 1) for camp in camps_names]\n",
" for kid in kids_names\n",
" ],\n",
" index=kids_names, columns=camps_names,\n",
")\n",
"variables"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>Cool Camp</th>\n",
" <th>Fun Camp</th>\n",
" <th>Super Camp</th>\n",
" <th>Sporty Camp</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>Ben</th>\n",
" <td>1</td>\n",
" <td>1</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>John</th>\n",
" <td>1</td>\n",
" <td>1</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Tom</th>\n",
" <td>1</td>\n",
" <td>1</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Chris</th>\n",
" <td>1</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Josh</th>\n",
" <td>1</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Alan</th>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>1</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Timmy</th>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>1</td>\n",
" <td>1</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Greg</th>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" <td>0</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" Cool Camp Fun Camp Super Camp Sporty Camp\n",
"Ben 1 1 0 0\n",
"John 1 1 0 0\n",
"Tom 1 1 0 0\n",
"Chris 1 0 0 0\n",
"Josh 1 0 0 0\n",
"Alan 0 0 1 1\n",
"Timmy 0 0 1 1\n",
"Greg 0 0 0 0"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"possible_assignments = pd.DataFrame(\n",
" [\n",
" [1 if kids_by_name[kid].sport in camps_by_name[camp].sports else 0\n",
" for camp in camps_names]\n",
" for kid in kids_names\n",
" ],\n",
" index=kids_names, columns=camps_names,\n",
")\n",
"possible_assignments"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"problem = pulp.LpProblem('Camp Assignment', sense=pulp.LpMaximize)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"# Objective\n",
"problem += variables.sum().sum(), 'Maximize number of kids going to camp'"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"# Only one camp per kid\n",
"for kid in kids_names:\n",
" problem += variables.loc[kid, :].sum() <= 1, f'Only 1 camp per kid ({kid})'\n",
" \n",
"# Limited number of spots per camp\n",
"for camp in camps_names:\n",
" spots = camps_by_name[camp].open_spots\n",
" problem += variables.loc[:, camp].sum() <= spots, f'Only {spots} spots in {camp}'\n",
" \n",
"# Only go if favorite sport is available\n",
"problem += possible_assignments.eq(0).multiply(variables).sum().sum() == 0, 'Only go if favorite sport is available'"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"1"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Solve the problem\n",
"problem.solve()"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>Cool Camp</th>\n",
" <th>Fun Camp</th>\n",
" <th>Super Camp</th>\n",
" <th>Sporty Camp</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>Ben</th>\n",
" <td>True</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" </tr>\n",
" <tr>\n",
" <th>John</th>\n",
" <td>False</td>\n",
" <td>True</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Tom</th>\n",
" <td>False</td>\n",
" <td>True</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Chris</th>\n",
" <td>True</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Josh</th>\n",
" <td>True</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Alan</th>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" <td>True</td>\n",
" <td>False</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Timmy</th>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" <td>True</td>\n",
" </tr>\n",
" <tr>\n",
" <th>Greg</th>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" <td>False</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" Cool Camp Fun Camp Super Camp Sporty Camp\n",
"Ben True False False False\n",
"John False True False False\n",
"Tom False True False False\n",
"Chris True False False False\n",
"Josh True False False False\n",
"Alan False False True False\n",
"Timmy False False False True\n",
"Greg False False False False"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Extract solution\n",
"assignments = variables.applymap(lambda x: x.value()).astype(bool)\n",
"assignments"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'Ben': 'Cool Camp',\n",
" 'John': 'Fun Camp',\n",
" 'Tom': 'Fun Camp',\n",
" 'Chris': 'Cool Camp',\n",
" 'Josh': 'Cool Camp',\n",
" 'Alan': 'Super Camp',\n",
" 'Timmy': 'Sporty Camp',\n",
" 'Greg': None}"
]
},
"execution_count": 12,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Get solution to preferred format\n",
"answer = {\n",
" kid: next((camp for camp, is_assigned in row.items() if is_assigned), None)\n",
" for kid, row in assignments.iterrows() \n",
"}\n",
"answer"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"True"
]
},
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"answer == expected_answer"
]
}
],
"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.7.3"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment