stable_marriage.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "from random import sample\n",
    "import pprint\n",
    "import csv\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "## SET THESE VARIABLES\n",
    "## Ideally, num_teams * max_team_size = num_devs\n",
    "## Otherwise, num_teams * (max_team_size - 1) < num_devs <= num_teams * max_team_size, also works\n",
    "num_teams = 4\n",
    "num_devs = 22\n",
    "max_team_size = 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "def stable_marriage(teams_rankings, devs_rankings):\n",
    "    '''\n",
    "    Performs stable matching algorithm where each team can propose to up to `num_teams` number of devs\n",
    "    \n",
    "    Parameters\n",
    "    -----------\n",
    "    teams_rankings: dictionary\n",
    "        key is team name, value is list of devs where lower index is higher preference\n",
    "    devs_rankings: dictionary\n",
    "        key is dev name, value is list of teams where lower index is higher preference\n",
    "    \n",
    "    Returns\n",
    "    -----------\n",
    "    dev_matching: dictionary\n",
    "        key is dev, value is the team dev is matched to\n",
    "    '''\n",
    "    \n",
    "    teams = teams_rankings.keys()\n",
    "    devs = devs_rankings.keys()\n",
    "    \n",
    "    matching = [{team: []} for team in teams]\n",
    "    team_availability = dict(zip(teams, [\"free\"] * len(teams)))\n",
    "    team_size = dict(zip(teams, [0] * len(teams)))\n",
    "    dev_availability = dict(zip(devs, [\"free\"] * len(devs)))\n",
    "    team_counter = dict(zip(teams, [0] * len(teams)))\n",
    "    dev_matching = {}\n",
    "    \n",
    "    while not teams_full(team_size):\n",
    "        if proposals_completed(team_counter):\n",
    "            break\n",
    "        \n",
    "        for team in teams:\n",
    "            if team_availability[team] != \"free\" or team_counter[team] >= num_devs:\n",
    "                continue\n",
    "\n",
    "            dev_index = team_counter[team]\n",
    "            dev = teams_rankings[team][dev_index]\n",
    "\n",
    "            if dev_availability[dev] == \"free\":\n",
    "                dev_availability[dev] = \"engaged\"\n",
    "                dev_matching[dev] = team\n",
    "                team_size[team] += 1\n",
    "                if team_size[team] >= max_team_size:\n",
    "                    team_availability[team] = \"engaged\"\n",
    "\n",
    "            elif dev_availability[dev] == \"engaged\":\n",
    "                matched_team = dev_matching[dev]\n",
    "                proposed_team = team\n",
    "                dev_rankings = devs_rankings[dev]\n",
    "                if dev_rankings.index(proposed_team) < dev_rankings.index(matched_team):\n",
    "                    team_size[matched_team] -= 1\n",
    "                    if team_availability[matched_team] == \"engaged\":\n",
    "                        team_availability[matched_team] = \"free\"\n",
    "\n",
    "                    team_size[proposed_team] += 1\n",
    "                    if team_size[proposed_team] >= max_team_size:\n",
    "                        team_availability[proposed_team] = \"engaged\"\n",
    "\n",
    "                    dev_matching[dev] = proposed_team\n",
    "\n",
    "            team_counter[team] += 1\n",
    "    return dev_matching\n",
    "\n",
    "def teams_full(team_size):\n",
    "    '''\n",
    "    Returns true if sum of all team sizes is num_devs\n",
    "    \n",
    "    Parameters\n",
    "    -----------\n",
    "    team_size: dictionary\n",
    "        key is team, value is size of team\n",
    "\n",
    "    Returns\n",
    "    -----------\n",
    "    boolean\n",
    "    '''\n",
    "    devs_matched = 0\n",
    "    for _, size in team_size.items():\n",
    "        devs_matched += size\n",
    "    return devs_matched == num_devs\n",
    "\n",
    "def proposals_completed(team_counter):\n",
    "    '''\n",
    "    Returns true if all teams have proposed to all devs\n",
    "    \n",
    "    Parameters\n",
    "    -----------\n",
    "    team_counter: dictionary\n",
    "        key is team, value is number of devs team has proposed to\n",
    "\n",
    "    Returns\n",
    "    -----------\n",
    "    boolean\n",
    "    '''\n",
    "    for _, count in team_counter.items():\n",
    "        if count < num_devs - 1:\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "def dev_to_team_matching(dev_matching):\n",
    "    '''\n",
    "    Converts dev-to-team matching to team-to-dev matching\n",
    "    \n",
    "    Parameters\n",
    "    -----------\n",
    "    dev_matching: dictionary\n",
    "        key is dev, value is the team dev is matched to\n",
    "\n",
    "    Returns\n",
    "    -----------\n",
    "    team_matching: dictionary\n",
    "        key is team, value is list containing devs matched to that team\n",
    "    '''\n",
    "    team_matching = {}\n",
    "    for dev, team in dev_matching.items():\n",
    "        if team in team_matching.keys():\n",
    "            team_matching[team].append(dev)\n",
    "        else:\n",
    "            team_matching[team] = [dev]\n",
    "    return team_matching\n",
    "\n",
    "def get_team_names():\n",
    "    '''\n",
    "    Gets team names from team_ranking.csv\n",
    "    \n",
    "    Returns\n",
    "    -----------\n",
    "    teams: list\n",
    "        list of team names\n",
    "    '''\n",
    "    with open(\"team_ranking.csv\", \"r\") as f:\n",
    "        reader = csv.reader(f, delimiter=\"\\t\")\n",
    "\n",
    "        for i, line in enumerate(reader):\n",
    "            return line[0].split(',')[1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Atlassian', 'Hackerrank', 'Netskope', 'Pantheon']\n"
     ]
    }
   ],
   "source": [
    "team_names = get_team_names()\n",
    "print(team_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Andi Gu': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Andrew Lieu': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Bianca Lee': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Brian DeLeonardis': ['Netskope', 'Pantheon', 'Hackerrank', 'Atlassian'],\n",
      " 'Carolyn Duan': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Cody Zeng': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Daniel Grimshaw': ['Hackerrank', 'Atlassian', 'Netskope', 'Pantheon'],\n",
      " 'Erica Kong': ['Hackerrank', 'Atlassian', 'Netskope', 'Pantheon'],\n",
      " 'Forest Hu': ['Netskope', 'Pantheon', 'Hackerrank', 'Atlassian'],\n",
      " 'Isabelle Zhou': ['Hackerrank', 'Atlassian', 'Netskope', 'Pantheon'],\n",
      " 'Ishan Vacchani': ['Netskope', 'Pantheon', 'Hackerrank', 'Atlassian'],\n",
      " 'Jared Rosner': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Kelvin Jue': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Kush Khanolkar': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Kyle Tse': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Lang Tran': ['Netskope', 'Pantheon', 'Hackerrank', 'Atlassian'],\n",
      " 'Matthew Soh': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Nilay Khatore': ['Hackerrank', 'Atlassian', 'Netskope', 'Pantheon'],\n",
      " 'Richard Liu': ['Atlassian', 'Netskope', 'Pantheon', 'Hackerrank'],\n",
      " 'Steven Tsay': ['Netskope', 'Pantheon', 'Hackerrank', 'Atlassian'],\n",
      " 'Trevor Aquino': ['Hackerrank', 'Atlassian', 'Netskope', 'Pantheon'],\n",
      " 'William Zhou': ['Netskope', 'Pantheon', 'Hackerrank', 'Atlassian']}\n"
     ]
    }
   ],
   "source": [
    "devs_ranking = {}\n",
    "\n",
    "with open(\"dev_ranking.csv\", \"r\") as f:\n",
    "    reader = csv.reader(f, delimiter=\"\\t\")\n",
    "    for i, line in enumerate(reader):\n",
    "        if i == 0:\n",
    "            continue\n",
    "        args = line[0].split(',')\n",
    "        dev_name = args[0]\n",
    "        dev_ranking = args[1:]\n",
    "            \n",
    "        devs_ranking[dev_name] = dev_ranking\n",
    "        \n",
    "pprint.pprint(devs_ranking)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Atlassian': ['Nilay Khatore',\n",
      "               'Trevor Aquino',\n",
      "               'Erica Kong',\n",
      "               'Daniel Grimshaw',\n",
      "               'Isabelle Zhou',\n",
      "               'Kelvin Jue',\n",
      "               'Carolyn Duan',\n",
      "               'Andrew Lieu',\n",
      "               'Bianca Lee',\n",
      "               'Andi Gu',\n",
      "               'Richard Liu',\n",
      "               'William Zhou',\n",
      "               'Forest Hu',\n",
      "               'Lang Tran',\n",
      "               'Steven Tsay',\n",
      "               'Ishan Vacchani',\n",
      "               'Brian DeLeonardis',\n",
      "               'Kyle Tse',\n",
      "               'Cody Zeng',\n",
      "               'Kush Khanolkar',\n",
      "               'Jared Rosner',\n",
      "               'Matthew Soh'],\n",
      " 'Hackerrank': ['Jared Rosner',\n",
      "                'Matthew Soh',\n",
      "                'Nilay Khatore',\n",
      "                'Trevor Aquino',\n",
      "                'Erica Kong',\n",
      "                'Daniel Grimshaw',\n",
      "                'Isabelle Zhou',\n",
      "                'Kelvin Jue',\n",
      "                'Carolyn Duan',\n",
      "                'Andrew Lieu',\n",
      "                'Bianca Lee',\n",
      "                'Andi Gu',\n",
      "                'Richard Liu',\n",
      "                'William Zhou',\n",
      "                'Forest Hu',\n",
      "                'Lang Tran',\n",
      "                'Steven Tsay',\n",
      "                'Ishan Vacchani',\n",
      "                'Brian DeLeonardis',\n",
      "                'Kyle Tse',\n",
      "                'Cody Zeng',\n",
      "                'Kush Khanolkar'],\n",
      " 'Netskope': ['Cody Zeng',\n",
      "              'Kush Khanolkar',\n",
      "              'Jared Rosner',\n",
      "              'Matthew Soh',\n",
      "              'Nilay Khatore',\n",
      "              'Trevor Aquino',\n",
      "              'Erica Kong',\n",
      "              'Daniel Grimshaw',\n",
      "              'Isabelle Zhou',\n",
      "              'Kelvin Jue',\n",
      "              'Carolyn Duan',\n",
      "              'Andrew Lieu',\n",
      "              'Bianca Lee',\n",
      "              'Andi Gu',\n",
      "              'Richard Liu',\n",
      "              'William Zhou',\n",
      "              'Forest Hu',\n",
      "              'Lang Tran',\n",
      "              'Steven Tsay',\n",
      "              'Ishan Vacchani',\n",
      "              'Brian DeLeonardis',\n",
      "              'Kyle Tse'],\n",
      " 'Pantheon': ['Brian DeLeonardis',\n",
      "              'Kyle Tse',\n",
      "              'Cody Zeng',\n",
      "              'Kush Khanolkar',\n",
      "              'Jared Rosner',\n",
      "              'Matthew Soh',\n",
      "              'Nilay Khatore',\n",
      "              'Trevor Aquino',\n",
      "              'Erica Kong',\n",
      "              'Daniel Grimshaw',\n",
      "              'Isabelle Zhou',\n",
      "              'Kelvin Jue',\n",
      "              'Carolyn Duan',\n",
      "              'Andrew Lieu',\n",
      "              'Bianca Lee',\n",
      "              'Andi Gu',\n",
      "              'Richard Liu',\n",
      "              'William Zhou',\n",
      "              'Forest Hu',\n",
      "              'Lang Tran',\n",
      "              'Steven Tsay',\n",
      "              'Ishan Vacchani']}\n"
     ]
    }
   ],
   "source": [
    "dev_list = [[-1] * num_devs for i in range(num_teams)]\n",
    "teams_ranking = dict(zip(team_names, dev_list))\n",
    "with open(\"team_ranking.csv\", \"r\") as f:\n",
    "    reader = csv.reader(f, delimiter=\"\\t\")\n",
    "    \n",
    "    for i, line in enumerate(reader):\n",
    "        if i == 0:\n",
    "            continue\n",
    "        args = line[0].split(',')\n",
    "        dev_name = args[0]\n",
    "        dev_rankings = args[1:]\n",
    "        \n",
    "        for team, ranking in zip(team_names, dev_rankings):\n",
    "            teams_ranking[team][int(ranking) - 1] = dev_name\n",
    "            \n",
    "pprint.pprint(teams_ranking)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [],
   "source": [
    "stable_matching = stable_marriage(teams_ranking, devs_ranking)\n",
    "team_matching = dev_to_team_matching(stable_matching)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Atlassian': ['Carolyn Duan',\n",
       "  'Kelvin Jue',\n",
       "  'Bianca Lee',\n",
       "  'Richard Liu',\n",
       "  'Andi Gu',\n",
       "  'Andrew Lieu'],\n",
       " 'Hackerrank': ['Daniel Grimshaw',\n",
       "  'Trevor Aquino',\n",
       "  'Nilay Khatore',\n",
       "  'Isabelle Zhou',\n",
       "  'Erica Kong',\n",
       "  'Ishan Vacchani'],\n",
       " 'Netskope': ['William Zhou',\n",
       "  'Kush Khanolkar',\n",
       "  'Cody Zeng',\n",
       "  'Matthew Soh',\n",
       "  'Forest Hu',\n",
       "  'Jared Rosner'],\n",
       " 'Pantheon': ['Steven Tsay', 'Lang Tran', 'Brian DeLeonardis', 'Kyle Tse']}"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "team_matching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"stable_matching.txt\", 'w') as f:\n",
    "    for key, value in team_matching.items():\n",
    "        f.write('%s:%s\\n' % (key, value))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}