HungarianAlgorithm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hungarian Algorithm\n",
    "\n",
    "James Yu, 7 March 2021, Updated 31 March 2021\n",
    "\n",
    "## Intro\n",
    "\n",
    "This notebook implements the Hungarian Algorithm in Python as outlined in [this reading](https://montoya.econ.ubc.ca/Econ514/hungarian.pdf) by Dr. Michael Peters.\n",
    "\n",
    "It upgrades a graph-based implementation of the algorithm shown in [these notes](http://www.cse.ust.hk/~golin/COMP572/Notes/Matching.pdf) to use matrices and explicitly define the form of the alternating tree.\n",
    "\n",
    "For a summary and full code, skip to the **Recap** section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Hungarian Algorithm is a type of matching algorithm capable of finding a subgame-perfect Nash equilibrium to a strategic game where one set of agents are selecting a choice from another set of agents. For example, one may have a group of workers applying to firms, or students applying to universities. The algorithm will find the particular matching of agents which maximizes the total surplus of the system; that is, the sum of earnings from all agents in the outcome will be the maximum possible value."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dr. Peters' notes provides a simple manual example so lets use it to develop our algorithm implementation:\n",
    "\n",
    "$$\n",
    "\\begin{array}{|c|c|c|}\n",
    "\\hline \\textbf{} & \\textbf{Firm A} & \\textbf{Firm B}\\\\\\hline\n",
    "  \\textbf{Worker 1} & 100 & 50 \\\\\\hline\n",
    "  \\textbf{Worker 2} & 80 & 26 \\\\\\hline\n",
    "\\end{array}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, the numbers in the boxes correspond to revenues received by the firms for employing a particular worker. We wish to find the matching between workers and firms which maximizes total revenue. \n",
    "\n",
    "Such a matching should have mutually exclusive pairings; in other words, no two firms matched to the same worker and no two workers matched to the same firm."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Put Payoffs into Numpy Matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This part just initializes our problem. There's no theory here. We use NumPy to store the matrix of revenues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[100,  50],\n",
       "       [ 80,  26]])"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "matrx = np.array([\n",
    "    [100, 50], \n",
    "    [80, 26]])\n",
    "\n",
    "matrx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above gives us a simple \"list of lists\" which we can treat as a matrix. To access a specific element, we can use the following notation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "50"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "matrx[0, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first number, 0, corresponds to the row index; and the second number, 1, corresponds to the column index.\n",
    "\n",
    "We could also retrieve an entire row or column as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([100,  80])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "matrx[:, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([100,  50])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "matrx[0, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Determine an Initial Set of Potentials"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we can manipulate matrices, we can start building our algorithm. The first step is to initialize **potentials** which are values representing \"potential\" outcomes for a distribution of wages and profits for our revenues.\n",
    "\n",
    "Each worker has a potential corresponding to the wages they could get, and each firm has a potential corresponding to the profits they could get. Recall profit is revenue minus cost, so there is intuition here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On paper this looks like:\n",
    "\n",
    "$$\n",
    "\\begin{array}{|c|c|c|c|}\n",
    "\\hline \\textbf{} & \\textbf{Potentials} & \\textbf{Firm A} & \\textbf{Firm B}\\\\\\hline\n",
    "  \\textbf{Potentials} & \\textbf{} & 0 & 0 \\\\\\hline\n",
    "  \\textbf{Worker 1} & 100 & 100 & 50 \\\\\\hline\n",
    "  \\textbf{Worker 2} & 80 & 80 & 26 \\\\\\hline\n",
    "\\end{array}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start, we set the potentials for the workers to be the maximum revenue seen in each row (i.e. the maximum potential wages one could possibly get).\n",
    "\n",
    "We set potentials for firms to be zero which balances this out. The sum of row and column potentials must be at least as good as revenue for the potentials to be valid; if not, we have a problem. The readings explain this in more detail, so we will take this as given."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In code, we need an auxiliary method of storing the potentials, since we'll be manipulating them somewhat separately from the revenues. Revenues never change, since this is static data, but potentials change as our algorithm improves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[100, 80]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "size = 2\n",
    "\n",
    "rpotentials = []\n",
    "cpotentials = [0 for i in range(size)]\n",
    "for x in range(len(matrx)):\n",
    "    rpotentials.append(max(matrx[x, :]))\n",
    "\n",
    "rpotentials"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code sets the column (firm) potentials to be all zero, and sets the row potentials to be the maximum value present in each row."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2a: Equality Graph\n",
    "\n",
    "This step has no code associated with it. Recall an **equality graph** is a set of worker-firm pairs such that their potentials add up to the revenue. For example, if a worker had potential 50 and a firm had potential 0 and revenue 50, 50+0=50 and we would have this in the equality graph. We'll use this concept in the next step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: From Alternating Path to Augmenting Path\n",
    "\n",
    "This is where we slightly deviate from what is explained in the readings. The explanation given in the readings provides the same thing, but this is the hidden mathematics underneath.\n",
    "\n",
    "First, note that what we are looking for is a **matching**, or a set of worker-firm pairs that form the Nash equilibrium pairings we need (where everyone gets their own pair). \n",
    "\n",
    "The only pairs that are capable of being included are those also present in the equality graph, which is why that concept is so important.\n",
    "\n",
    "Second, note that as the algorithm progresses, we may end up changing the worker-firm pairs we included in the matching. Whenever we refer to a matching in the explanation that follows, we refer to the matching we get from the current state of our algorithm.\n",
    "\n",
    "As in the readings, an **alternating path** is a sequence of worker-firm pairs such that two adjacent pairs either have a shared row (worker) or shared column (firm), and that the pairs alternate between being included in our current matching result and being excluded. \n",
    "\n",
    "In other words, if one pair in the sequence is in our match, the next won't be, but the one after that will, and etc.\n",
    "\n",
    "What the notes do not describe is the fact that the alternating path, at its core, is not a path between pairs but a path between individuals.\n",
    "\n",
    "Here's a clearer explanation. Start with a particular worker/row that we **haven't matched yet**. This will always be the starting point for our alternating path.\n",
    "\n",
    "$$W_1$$\n",
    "\n",
    "Here alternating also happens to refer to the fact that our path alternates between workers and firms. Given that we started with a worker, we want to find a firm to add to our path.\n",
    "\n",
    "If we're lucky, we can find a firm such that the worker-firm pair is in the equality graph. We might have something like:\n",
    "\n",
    "$$W_1 \\textbf{->} F_A$$\n",
    "\n",
    "(Bold arrow represents that a worker-firm pair is in our match).\n",
    "\n",
    "Now, in order to add one of these pairs to our matching, we need to make sure neither is already matched. If that's true (i.e. **if the firm is unmatched**), we're **done with the path** and can move on to another worker.\n",
    "\n",
    "But **if not**, here's where things get interesting. We know for sure that $W_1$ is unmatched because we said so. This means $F_A$ is matched to some worker, let's say $W_2$. So we add it to our path. Starting from scratch:\n",
    "$$W_1 -> F_A$$\n",
    "\n",
    "becomes\n",
    "\n",
    "$$W_1 -> F_A \\textbf{->} W_2$$\n",
    "\n",
    "**(Note the bolded arrow representing a match)**. We still need to match $W_1$ to *something*, so we keep going. The idea is to come up with enough matches to match everyone.\n",
    "\n",
    "Perhaps we have some other firm $F_B$ that we haven't matched yet. If we wanted, we could adjust the potentials to put $F_B$, $W_2$ in our equality graph. This seems rather pointless ($W_2$ is matched already!) until you realize the following:\n",
    "\n",
    "$$W_1 -> F_A \\textbf{->} W_2 -> F_B$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A key characteristic about an alternating path is that every pairing in it is in the equality graph. That is, any worker-firm pair we can find in this pathway could *potentially* be a match. We've currently configured it so there happens to be just one match in a set of three pairings.\n",
    "\n",
    "What we can do, however, is a trick that allows us to switch the pairs we want in our match. If we have a path like this where the first entry is an unmatched worker and the last entry is an unmatched firm, we can invert every pair such that matches become unmatched and non-matched become matched. This gives us:\n",
    "\n",
    "$$W_1 \\textbf{->} F_A -> W_2 \\textbf{->} F_B$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Two individual matches for four agents means everyone is matched and everyone is happy, so we'd be done.\n",
    "\n",
    "Any such path where the first and last entries are unmatched and we can invert the matches is called an **augmenting path**, because it lets us **augment** our matches by having one more than before."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The part that might be confusing is how we choose what firm to match to what worker. For now, let's not worry about this, and just assume we have some way to create this path. We'll explain later.\n",
    "\n",
    "Now we have to create structures in code to store these paths. We need a way to store the matched workers and the matched firms. We also need to know what's matched to what (the arrows)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "matching = []\n",
    "S = {0}\n",
    "T = set()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we have a list `matching` for the pairs we want to include in our match, a set `S` for our workers in the path, and a set `T` for the firms in our path. \n",
    "\n",
    "As above we want to start with some unmatched worker. It doesn't matter which one we pick, so we picked the first one, 0, and put it in our list of workers in the path."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Alternating Tree\n",
    "\n",
    "Here's where things get really interesting. And perhaps confusing. Let's tackle it step by step.\n",
    "\n",
    "In the readings, you may have seen something about **degenerate paths** that you find alongside a particular alternating path. There's a reason for this.\n",
    "\n",
    "What we *really* have is not just a *path*. In reality, we are creating a structure called a **tree**.\n",
    "\n",
    "You might be hearing this for the first time so we won't go into too much detail. For more information on trees, check out graph theory.\n",
    "\n",
    "In simpler terms, a tree is a just a bunch of connected lines. That's it. Its just like our path, but the path can branch out in multiple directions instead of just going in one direction. \n",
    "\n",
    "An **alternating tree** is one such group of lines where any pathway you trace is an alternating path. This might seem hard to visualize so here's an example.\n",
    "\n",
    "Let's say our alternating path was something like \n",
    "\n",
    "$$W_1 -> F_A \\textbf{->} W_2 -> F_B \\textbf{->} W_3$$\n",
    "\n",
    "(Just an example). Now, recall that our ultimate goal is to match workers with firms, and that we had two options. Either we found an unmatched firm that we can add to finish the path, or we added a firm that was already matched (along with its worker) to the path.\n",
    "\n",
    "We made a strange assumption that we could just stuff the firm to the end of our path and be done. This is actually not the case. In fact, it may be possible to stick it to *other* points in the path (as long as they are workers like the end of the path).\n",
    "\n",
    "Suppose we have a $F_C$ matched to some $W_4$, and suppose we did some magic that puts ($W_2$, $F_C$) in the equality graph. Again this seems like confusing magic now, but all will be revealed later.\n",
    "\n",
    "Then we can do something like this:\n",
    "\n",
    "$$W_1 -> F_A \\textbf{->} W_2 -> F_B \\textbf{->} W_3$$\n",
    ".\n",
    "$$\\downarrow$$\n",
    ".\n",
    "$$F_C \\textbf{->} W_4$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hopefully that rendered properly but there should be an arrow between $W_2$ and $F_C$. The point is, this is what an alternating tree might look like. We have two paths we could take, one to $W_3$ and one to $W_4$, and all of them are alternating between a non-match and match.\n",
    "\n",
    "At the end, when we find a free firm to finish the path, we won't flip all the matches in the tree like we did with our augmenting path from before. Instead, we just flip the matches in the single path connected to that firm. For example, if we attached a new firm $F_D$ to $W_4$ above, we'd only flip matches in the path that traces back to $W_1$. We don't touch the path leading to $W_3$.\n",
    "\n",
    "In general, where we add firms to our alternating tree depends on what the equality graph looks like, as we can only add pairs that are in the equality graph. We'll see this tree arises because of how we adjust the potentials, but for now let's just prep for this tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Node:\n",
    "    def __init__(self, val, parent = None):\n",
    "        self.val = val\n",
    "        self.parent = parent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Before we continue**, this code may seem unfamiliar. Here's all you need to know.\n",
    "\n",
    "Essentially we are defining a new type of variable called `Node`. `__init__` specifies what happens when we create one of these variables, and in this case we are giving the variable a few sub-variables. You can ignore `self`, which is just a name by which we can modify our variable.\n",
    "\n",
    "Here we give each `Node` a `val` and `parent`, which we'll explain below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now back to the theory. The firms and workers in our alternating tree are little points that we connect together. Some literatures call these points vertices, while others call them nodes. Here we call them nodes.\n",
    "\n",
    "A **node**, which can be a worker or firm, has both a name `val` (the index of the worker/firm in our matrix) and a possible `parent` node.\n",
    "\n",
    "A **parent node** is the node that points to this one. For example, $W_4$ in our simple alternating tree has parent $F_C$, and $F_A$ has parent $W_1$. $W_1$ has no parent because its the first node, sometimes called the **root node**. This is why `parent` has a **default value** of `None` in our `__init__` function, because it might be the case that a node has no parent.\n",
    "\n",
    "This means we're actually storing our trees backwards. Instead of making connections left to right like in our diagram, we store our nodes pointing right to left (from node to parent). If you look very very carefully at our example tree, you'll notice every node has exactly one parent. This is not a coincidence, and this allows us to always get back to the root node (front of the path) from any node.\n",
    "\n",
    "On the other hand, if you go left to right, some nodes branch out in multiple directions. This is harder to keep track of in code because you could have one branch, or two branches, or many branches. So we store pathways in reverse.\n",
    "\n",
    "The code up there just tells us what one of these nodes looks like; we have to initialize our alternating tree in the program too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree_root = Node(0)\n",
    "x_nodes = {0: tree_root}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we make one node to start with, which comes all the way from our `S` set earlier. It contains the first worker in our payoff matrix, which is the one we want to try to match first.\n",
    "\n",
    "We also have a **dictionary** `x_nodes` which is essentially a list of workers which we can access by giving the worker IDs. For example, since our first worker is worker 0, we can do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.Node at 0x7f02dfe42190>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_nodes[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "to get the node associated with that worker. This seems like a cryptic ID, but we can unpack a node as follows. Recall every node has a `val` and `parent`, so we can do something like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_nodes[0].val"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "to get the `val` associated with the node. We can do the same for `parent` if we want."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: How do we match? Neighbours.\n",
    "\n",
    "We're done initializing now, so we can start making the mechanics of the algorithm. In order to reach a matching, we have to make matches.\n",
    "\n",
    "To do this, recall matches must come from the equality graph, so we need a way to find pairings that are in the equality graph.\n",
    "\n",
    "We're starting from worker 0, so we want to find some firm that might be in the equality graph when paired with worker 0.\n",
    "\n",
    "This introduces the concept of **neighbours**. A **neighbour** of a worker is a firm such that the worker-firm pair is in the equality graph. If you've forgotten by now, this means that the potential of the worker plus the potential of the firm is the revenue of the firm.\n",
    "\n",
    "We therefore need a function to sort through the payoff matrix to find them. We do the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def neighbours(wset):\n",
    "    # find all firms in equality graph\n",
    "    result = []\n",
    "    for x in wset:\n",
    "        # get row of firms for worker x\n",
    "        nbs = matrx[x, :]\n",
    "        for y in range(len(nbs)):\n",
    "            if nbs[y] == rpotentials[x] + cpotentials[y]:\n",
    "                result.append([x, y])\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a bit to unpack. First of all, `wset` is a list of workers we want to find the neighbours of. For example, if we want the neighbours for worker 0, we pass in `[0]` as wset.\n",
    "\n",
    "Next, we run `for x in wset` to iterate over the workers `x` that we need. For each of these, we get the row in the payoff matrix corresponding to that worker by calling `matrx[x, :]` as before.\n",
    "\n",
    "Using this row, we check each element to see if there are any firms which, when paired with this worker, end up in the equality graph.\n",
    "\n",
    "This amounts to checking every firm `y` in the row and seeing if the potential for worker `x` plus the potential for firm `y` is equal to the revenue payoff in the row. If so, we append our `x`, `y` pair to `result`. To end, we return it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus, we can now get neighbouring firms of any worker which fall into the equality graph.\n",
    "\n",
    "The point of this was that we needed to match our worker 0 to some firm, so lets run `neighbours` to find a match. For our algorithm, we'll run `neighbours` on our set `S` of workers in our alternating tree, because remember we can attach a firm to any worker in the tree if it turns out that worker-firm pair is in the equality graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0, 0]]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NS = neighbours(S)\n",
    "NS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It looks like `[0, 0]` is a worker-firm pair. If we want to see the associated revenue, we can do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "matrx[0, 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So the top left entry with revenue 100 is in the equality graph. This is promising, and firm 0 isn't matched yet, so let's add it to our match. Remember that when we reach an unmatched firm, we can just stop and ditch the path/tree and add the match to our matches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "matching.append([0, 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is good. Now we keep going. We want to find the next worker to be matched, so lets pull one from the workers we have left."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "free = list(set(range(size)) - set([m[0] for m in matching]))\n",
    "free"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a bit more to unpack. Here we create a set from `range(size)`, which is the total list of workers we have since we numbered them numerically. `size` is 2 so `range(size)` is 0 and 1.\n",
    "\n",
    "The other set takes `m[0]` for `m` in `matching`. This means we take the first element of every matched worker-firm pair. Since the first element is the worker, or x, this is a list of matched workers.\n",
    "\n",
    "When we do `set(range(size)) - set([m[0] for m in matching])`, recall that a set is a collection of unique objects. So when we subtract the matched workers from the total workers, we get a set of unmatched workers.\n",
    "\n",
    "Finally, we convert it to a list so we can get an index out. The result is that `free` is `[1]`, which contains worker 1, which we know is our only unmatched worker, so this works."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Whenever we reach an augmenting path in an alternating tree, recall we mentioned we were \"done\" with the path. We really do mean done; when we continue the algorithm, we wipe the tree we already have and start over. Every worker gets their own alternating tree as every worker has their own matches that need to be found. Thus, lets wipe the tree and add our new worker instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree_root = Node(free[0])\n",
    "x_nodes = {free[0]: tree_root}\n",
    "S = {free[0]}\n",
    "T = set()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note we never actually used any of this to determine our match. Eventually we will, but in a smaller example, we terminate too quickly for the tree to gain any size.\n",
    "\n",
    "At this point, we have a new starting worker, so we can go back to the beginning of the algorithm and continue. This is the algorithmy part of the algorithm because we're just repeating the same procedure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[1, 0]]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NS = neighbours(S)\n",
    "NS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now here's the tricky part. If you look at this, we're pairing worker 1 with firm 0. But we already matched firm 0! So we can't add this to our matches and instead we must extend our alternating tree. Recall this is the second case of adding firms to our path."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In code, we can tell if we have this second case by doing this:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pair = [1, 0]\n",
    "pair[1] not in [m[1] for m in matching]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: If Match is Already Matched, Expand Tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`False` means that `pair[1]`, or the firm we want to match, is in fact in the list of `m[1]`s, or matched firms (since remember the second element `m[1]` of every match is the firm).\n",
    "\n",
    "So if this expression is false, we want to expand our alternating tree. We can do so like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "matching_x = next(m[0] for m in matching if m[1] == pair[1])\n",
    "S.add(matching_x)\n",
    "T.add(pair[1])\n",
    "source = x_nodes[pair[0]]\n",
    "y_node = Node(pair[1], source)\n",
    "x_node = Node(matching_x, y_node)\n",
    "x_nodes[matching_x] = x_node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{1: <__main__.Node at 0x7f02dfe54880>, 0: <__main__.Node at 0x7f02dfe5d8e0>}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "<__main__.Node object at 0x7f02dfe5d940>\n",
      "0\n",
      "<__main__.Node object at 0x7f02dfe54880>\n",
      "1\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "print(x_nodes[0].val)\n",
    "print(x_nodes[0].parent)\n",
    "print(x_nodes[0].parent.val)\n",
    "print(x_nodes[0].parent.parent)\n",
    "print(x_nodes[0].parent.parent.val)\n",
    "print(x_nodes[0].parent.parent.parent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's unpack what we just did.\n",
    "\n",
    "`next(m[0] for m in matching if m[1] == pair[1])` is a line of code which finds the worker matched to the firm we found out was already matched. Specifically, we return the `next` instance of worker `m[0]` if firm `m[1]` is the `pair[1]` firm we *wanted* to match from before.\n",
    "\n",
    "So now we have three agents. We have the previously matched worker-firm pair, and a worker that didn't get matched but is still in the equality graph with the firm.\n",
    "\n",
    "We can therefore add this pre-matched worker-firm pair to our alternating tree by adding the worker to `S` and the firm to `T`.\n",
    "\n",
    "This puts them in our sets of nodes but we also need to save the order in which they go in our tree. First, we retrieve the node corresponding to the worker `pair[0]` that we failed to match.\n",
    "\n",
    "We create a new node for firm `pair[1]` and set it to have worker `pair[0]` as its parent. This now looks like:\n",
    "\n",
    "$$\\text{pair[0]} -> \\text{pair[1]}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We saw that firm `pair[1]` was already matched to a worker which we saved in `matching_x`, so we create a new node for that matching worker and set it to have parent `pair[1]`. This gives:\n",
    "\n",
    "$$\\text{pair[0]}->\\text{pair[1]}\\textbf{->}\\text{matching_x}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note in future iterations, `pair[0]` might end up being some intermediate node in our alternating tree as we mentioned before. Adding these extra nodes, however, will always create a new branch and/or extend an existing branch since we aren't connecting anything else to `matching_x` yet. In tree terminology, we call this endpoint a **leaf node**.\n",
    "\n",
    "Anyway, to end off, we add our new worker node to the indexed list/dict of worker nodes with `x_nodes[matching_x] = x_node`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might be wondering why we don't have a lookup table for the firm/y nodes. It turns out we won't actually ever need to look up firm nodes because we'll always be reading or writing the alternating tree using the worker nodes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Go Back to Step 5 Again\n",
    "\n",
    "Now that we enlarged our tree, we go back and recompute neighbours. Which you might think is pointless because expanding the tree doesn't change the equality graph. \n",
    "\n",
    "But remember: when we enlarge the alternating tree, we add new worker nodes to the tree by putting them in `S`. Which means new workers to compute the neighbours of. Observe:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0, 0], [1, 0]]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NS = neighbours(S)\n",
    "NS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have two neighbours relative to the workers in `S`. Recall these are the pairs with revenues 100 and 80, which are also the current row potentials. This is no coincidence; those row potentials being equal to the revenues when the column potentials are zero is why these both fall in the equality graph.\n",
    "\n",
    "Now, in order to declare one of these to be a match, we have to find an unmatched firm to pair up, because recall we can only end our alternating tree with an unmatched firm. Its possible the worker is already matched, but remember if our alternating tree has some matched workers in it, because we started with an unmatched worker at the very beginning, we can simply switch around the pairs we matched and the pairs we didn't which matches this new firm.\n",
    "\n",
    "We can check in code if such a \"free\" firm exists by doing this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "set([b[1] for b in NS]) == T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`b[1] for b in NS` will return the second element of every pair in the neighbour list, which is the firm. Recall that in our tree structure, `T` consists of the firms that are in our alternating tree. If we did add a firm to our tree that was already matched, it would end up in `T`. If it was free/unmatched, recall we didn't add and instead destroyed the tree. Thus, if we are at this point, `T` must only contain matched firms.\n",
    "\n",
    "So if the set of matched firms in our tree is the same as the set of firms in our neighbour pairs, we don't have any free firms to match into our alternating tree. This is bad, because we need a match to keep going."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 8: When out of Matches, Update Potentials\n",
    "\n",
    "Finally we get to the magical step of updating the potentials that has been hinted upon the entire time. When we have no matches left, we can update the potentials for each agent to both preserve existing pairs that are in the equality graph, and add new pairs to the equality graph.\n",
    "\n",
    "Here's how it works. Recall the potentials for the workers/rows correspond to the highest revenue entry in the row. This means the potentials may be too high to include some of the lower revenue values.\n",
    "\n",
    "We want to therefore go through each worker/row **in our alternating tree** and figure out the smallest change to any potential that would allow us to include a new firm **not in our alternating tree** to possibly pair with a worker in our equality graph.\n",
    "\n",
    "That is, we want to find a pair such that among all our tree-present workers and among all our non-tree-present firms, the difference between the sum of the worker/firm potentials and the firm revenue is minimized.\n",
    "\n",
    "If we have this difference, we will then subtract it from all our tree-present workers and add it to all our **tree-present** firms. Note that since we checked the **non-tree-present** firms, this means the original pair with the smallest difference will now be in the equality graph since we'd be removing the difference from the worker potential without adding anything to the firm potential.\n",
    "\n",
    "In our simple original example, we had:\n",
    "\n",
    "$$\n",
    "\\begin{array}{|c|c|c|c|}\n",
    "\\hline \\textbf{} & \\textbf{Potentials} & \\textbf{Firm A} & \\textbf{Firm B}\\\\\\hline\n",
    "  \\textbf{Potentials} & \\textbf{} & 0 & 0 \\\\\\hline\n",
    "  \\textbf{Worker 1} & 100 & 100 & 50 \\\\\\hline\n",
    "  \\textbf{Worker 2} & 80 & 80 & 26 \\\\\\hline\n",
    "\\end{array}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Right now worker 1 is matched to firm A, which is preventing us from doing anything because all the equality graph entries are also in firm A (100+0=100, 80+0=80).\n",
    "\n",
    "So let's find the smallest difference between potentials and revenue for workers in the alternating tree and firms not in the alternating tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0, 1}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0}"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So worker 0 and worker 1 ($W_1$ and $W_2$, respectively, in our example) are in our tree, and firm 0 (firm A) is in the tree meaning firm B is not.\n",
    "\n",
    "Let's do some math. For the first worker, 100+0-50 = 50. For the second, 80+0-26 = 54. 50 is less than 54, so we'll subtract 50.\n",
    "\n",
    "This gives:\n",
    "\n",
    "$$\n",
    "\\begin{array}{|c|c|c|c|}\n",
    "\\hline \\textbf{} & \\textbf{Potentials} & \\textbf{Firm A} & \\textbf{Firm B}\\\\\\hline\n",
    "  \\textbf{Potentials} & \\textbf{} & 0+50=50 & 0 \\\\\\hline\n",
    "  \\textbf{Worker 1} & 100-50=50 & 100 & 50 \\\\\\hline\n",
    "  \\textbf{Worker 2} & 80-50=30 & 80 & 26 \\\\\\hline\n",
    "\\end{array}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall we subtract from workers in the tree and add to firms in the tree. Now observe the entries in the equality graph:\n",
    "\n",
    "- 50+50=100\n",
    "- 30+50=80\n",
    "- 50+0=50\n",
    "\n",
    "By doing this computation, because the difference we used was that of some worker-firm pair, and we corrected the difference by adding it back to tree-present firms, we've maintained the existing equality graph while adding a new element. This is a very important conclusion: **every step of the algorithm, we either increase the equality graph, increase the tree, or add a new match**. Since the algorithm is continuously increasing the size of *something* in a finitely-sized set of agents, we will *eventually* hit some sort of endpoint.\n",
    "\n",
    "In code, we can compute the potential adjustment as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_potentials():\n",
    "    global rpotentials, cpotentials\n",
    "    big = np.inf\n",
    "    for dx in S:\n",
    "        for dy in set(range(size)) - set(T):\n",
    "            weight = matrx[dx, dy]\n",
    "            alpha = rpotentials[dx] + cpotentials[dy] - weight\n",
    "            if alpha < big:\n",
    "                big = alpha\n",
    "\n",
    "    for dx in S:\n",
    "        rpotentials[dx] -= big\n",
    "\n",
    "    for dy in T:\n",
    "        cpotentials[dy] += big"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's unpack this. First we set a maximum value `big` to compare our differences against. We then iterate over all tree-present workers in `S` and non-tree-present firms in the complement of `T` (that is, subtracting `T` from the list of all firms).\n",
    "\n",
    "Over these worker-firm pairs, we retrieve the sum of the potentials and subtract the revenue from the payoff matrix for that pair. This gives us our difference. If it is smaller than the smallest difference we found so far, set the smallest difference to be this one.\n",
    "\n",
    "By the end, `big` will contain the smallest difference we found in our pairs. We then do the step of subtracting it from all tree-present workers in `S` and tree-present firms in `T`, as before. This completes the updating step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[50, 30]\n",
      "[50, 0]\n"
     ]
    }
   ],
   "source": [
    "update_potentials()\n",
    "print(rpotentials)\n",
    "print(cpotentials)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is exactly the result we found by hand. The potentials for the workers are now 50 and 30, while the potentials for the firms are now 50 and 0. \n",
    "\n",
    "Because this increases the size of the equality graph by **adding a free firm to it**, it **must** be the case that we now have a free firm to match. Let's recompute neighbours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0, 0], [0, 1], [1, 0]]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NS = neighbours(S)\n",
    "NS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Indeed, there it is. If we want to pick out our new firm specifically, we can do this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 1]"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pair = next(n for n in NS if n[1] not in set(T))\n",
    "pair"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is to say that `pair` is the next instance of a pair `n` in the `NS` neighbour list that doesn't have firm `n[1]` in the set `T` of firms in the tree. In other words, its the next pair that has a firm not in the tree.\n",
    "\n",
    "This part is particularly important. Recall in the middle of our algorithm we trashed our old tree and created a new one. This means that some matches may be present in our current tree, but **not all of them**. We therefore need to make sure that, even though this seems like a free firm, it really is a free firm and isn't already matched from before.\n",
    "\n",
    "We can do that check like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pair[1] not in [m[1] for m in matching]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This checks if the firm of the pair is not in the list of firms in the list of existing actual matches. It isn't, so we're good to go."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 9: Flip the Augmenting Path\n",
    "\n",
    "Now things become really interesting. We want to add this pair `[0, 1]` to our matching. However, recall that worker 0 is already matched.\n",
    "\n",
    "Also recall that worker 0 is in our alternating tree. If we want to check, we can call:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{1: <__main__.Node at 0x7f02dfe54880>, 0: <__main__.Node at 0x7f02dfe5d8e0>}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which proves that worker 0 is there. What results is now very interesting. Because our tree started from an unmatched worker, and we have here an unmatched firm with a matched worker, we can do that match-unmatch flip of the path that leads from the first worker to the worker matched with our firm. By doing so, the outer worker and our firm will now be associated with matches, which increases the size of the set of matches by 1. It also, however, deletes existing matches and adds some new ones in doing so (recall some pre-matched pairs get unmatched, and pre-unmatched ones get matches).\n",
    "\n",
    "In code, we have to figure out which path of our tree to go through. It will be whichever path contains the worker in the pair we want to add (because we're keeping track of what was already in the tree, a worker will never be in the tree twice). Also recall the worker ended up in the tree in the first place because it was matched and we wanted to use it (back at step 6), so we know for sure it will be in here.\n",
    "\n",
    "Since we need to find the specific worker, we can use our lookup list/dict."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "source = x_nodes[pair[0]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we walk the path. We start with the firm we want to add, which is currently unmatched with this worker. We'll therefore add our pair as a new match."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "matching.append(pair)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What we want to do now is keep going along the path and alternate between unmatching and matching. Since we'll always start by matching, we can take advantage of this and pre-configure some code to walk the path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "matched = 1\n",
    "while source.parent != None:\n",
    "    above = source.parent\n",
    "    if matched:\n",
    "        # if previously matched, this should be removed from matching\n",
    "        matching.remove([source.val, above.val])\n",
    "    else:\n",
    "        # if previous was a remove, this is a match\n",
    "        matching.append([above.val, source.val])\n",
    "\n",
    "    matched = 1 - matched\n",
    "    source = above"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's unpack this. `matched` is what we call a boolean flag. It switches between 1 and 0 on every iteration and lets us select whether we want to match or unmatch a pair. By default, its set to 1, which says that the previous iteration (adding our new firm) was a match, so this one should be an unmatch.\n",
    "\n",
    "To walk the path, recall we configured the tree nodes to have parents. Thus, we can start at the worker we just matched and walk backwards until we get to the original unmatched worker at the very front of the tree (which has no parent). Thus, while our node has a parent, we get the parent and store it in `above`.\n",
    "\n",
    "If we want to unmatch, we have the first case. Our first pair we encounter, which is the worker we just matched, was previously matched to some firm. Thus, since `source` was the worker/x node, its parent must be a firm (because we only add things in alternating sequences of workers and firms).\n",
    "\n",
    "Thus, we remove the pair `[source.val, above.val]`, or `[worker, firm]`, from our matching list. We then set `source = above` so we can get the parent of the firm for the next round, and we flip the `matched` variable to say we just unmatched.\n",
    "\n",
    "If instead we wanted to match, we get the second case. Matches happen after unmatches, meaning we start where `source` is some firm from the last iteration. This means its parent will be a worker.\n",
    "\n",
    "We can therefore add `[above.val, source.val`, or `[worker, firm]`, to our matching list. We then do the same state-flipping code from the unmatch case.\n",
    "\n",
    "Eventually, after a bunch of matches and unmatches, we will reach the first node in our tree, which has no parent. Since we loop while a parent exists, once we hit this node, we are done with the match update.\n",
    "\n",
    "The very last match must be the one with the first node and some firm, as the first node we know is unmatched so will become matched. Thus, we have one more match than before. (If this is confusing, scroll back to the diagram with the match-flipping example)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 10: Are we Done?\n",
    "\n",
    "We are done once we have enough matches for every worker and firm we have. There's a theorem in the readings that proves that if we do have a big enough matching, it must be the optimal matching, but we won't elaborate on it here.\n",
    "\n",
    "Let's check our matching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[0, 1], [1, 0]]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "matching"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are two matches, and we have four agents, so we're done! Let's take a look at what we found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matching: [[0, 1], [1, 0]]\n",
      "Revenue: [50, 80]\n",
      "Wages: [50, 30]\n",
      "Profits: [50, 0]\n",
      "Total Revenue: 130\n"
     ]
    }
   ],
   "source": [
    "print(\"Matching:\", matching)\n",
    "print(\"Revenue:\", [matrx[m[0], m[1]] for m in matching])\n",
    "print(\"Wages:\", rpotentials)\n",
    "print(\"Profits:\", cpotentials)\n",
    "print(\"Total Revenue:\", sum([matrx[m[0], m[1]] for m in matching]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our match is the corner entries of worker 0 with Firm 1, and worker 1 with Firm 0. In our simple example, that's worker  1 with Firm B and worker 2 with Firm A.\n",
    "\n",
    "This gives revenue 50 from the first pair and revenue 80 from the second for a total sum of 130. Note the alternative was the other diagonal, or 100+26=126, so indeed this is the optimal solution.\n",
    "\n",
    "Wages are actually the row potentials. This is the intuition for the potentials: now they actually represent the earnings that workers get out of the total revenue amount. The first worker gets revenue 50 while the second gets 30.\n",
    "\n",
    "Finally, profits are the column potentials, or the profit share of revenue. The first firm gets 50 while the second gets nothing.\n",
    "\n",
    "To check: 50+0 = 50, 30+50 = 80.\n",
    "\n",
    "Thus, the algorithm is complete. You might be wondering now why this is the optimal solution when we matched a 100 revenue point before. Recall we matched worker 0 with Firm 1 *after* we did that matching. It turns out that when the potentials are updated, matching worker 0 with Firm 1 becomes optimal because its in the equality graph. Worker 2 is not in the equality graph with Firm 1 (30 > 26), so Firm 1 will definitely take worker 1.\n",
    "\n",
    "This deletes worker 1 from being able to match with Firm 0, so Firm 0 must match with worker 2. This explains our result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recap\n",
    "\n",
    "To summarize, here are the steps of the algorithm.\n",
    "\n",
    "1. Initialize the revenue matrix.\n",
    "2. Initialize trivial potentials: max revenue of each row, and zeroes for the columns.\n",
    "3. Initialize the alternating tree with an unmatched worker, both as node sets and as a physical tree.\n",
    "4. While we don't yet have enough matchings for everyone:\n",
    "    1. Take all workers in our tree and find the firms that are paired with them in the equality graph.\n",
    "    2. If all of those firms are already in the tree, update potentials so we can get a firm that isn't.\n",
    "    3. If our firm is actually matched outside the tree, stick it into the tree along with its worker and **go back to step A**.\n",
    "    4. If not, add this firm to the match list with whatever worker it matches with.\n",
    "    5. If we have other workers and firms in the tree, conduct a swap of matches and unmatches along the path associated with the worker we matched.\n",
    "    6. Destroy the existing tree and make a new one with the next unmatched worker, if one exists. **Go back to step A if this is true**. If not, everyone is matched, and we are **done**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 11: An Actual Algorithm\n",
    "\n",
    "The above code iterated through the steps of the algorithm without looping. We can encapsulate the loops we defined to create a single unit of code capable of accomplishing what we just outlined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "\n",
    "Hungarian Algorithm No. 5 by James Yuming Yu\n",
    "7 March 2021\n",
    "\n",
    "Named after the 5 different theoretical implementations consulted before a sufficient one was found to base off.\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "class Node:\n",
    "    \"\"\"A simple node for an alternating tree.\"\"\"\n",
    "    \n",
    "    def __init__(self, val, parent = None):\n",
    "        self.val = val\n",
    "        self.parent = parent\n",
    "\n",
    "def run_hungarian(matrx):\n",
    "    \"\"\"Runs the Hungarian Algorithm on a given matrix and returns the optimal matching with potentials.\"\"\"\n",
    "    \n",
    "    size = matrx.shape[0]\n",
    "    \n",
    "    # Step 2: Generate trivial potentials\n",
    "    rpotentials = []\n",
    "    cpotentials = [0 for i in range(size)]\n",
    "    for i in range(len(matrx)):\n",
    "        row = matrx[i]\n",
    "        rpotentials.append(max(row))\n",
    "\n",
    "    # Step 3: Initialize alternating tree\n",
    "    matching = []\n",
    "    S = {0}\n",
    "    T = set()\n",
    "\n",
    "    tree_root = Node(0)\n",
    "    x_nodes = {0: tree_root}\n",
    "    \n",
    "    # Create helper functions\n",
    "\n",
    "    def neighbours(wset):\n",
    "        \"\"\"Finds all firms in equality graph with workers in wset.\"\"\"\n",
    "    \n",
    "        result = []\n",
    "        for x in wset:\n",
    "            # get row of firms for worker x\n",
    "            nbs = matrx[x, :]\n",
    "            for y in range(len(nbs)):\n",
    "                # check for equality\n",
    "                if nbs[y] == rpotentials[x] + cpotentials[y]:\n",
    "                    result.append([x, y])\n",
    "\n",
    "        return result\n",
    "    \n",
    "\n",
    "    def update_potentials():\n",
    "        \"\"\"Find the smallest difference between treed workers and untreed firms \n",
    "            and use it to update potentials.\"\"\"\n",
    "        \n",
    "        # when using functions in functions, if modifying variables, call nonlocal\n",
    "        nonlocal rpotentials, cpotentials \n",
    "        big = np.inf\n",
    "        # iterate over relevant pairs\n",
    "        for dx in S:\n",
    "            for dy in set(range(size)) - T:\n",
    "                # find the difference and check if its smaller than any we found before\n",
    "                weight = matrx[dx, dy]\n",
    "                alpha = rpotentials[dx] + cpotentials[dy] - weight\n",
    "                if alpha < big:\n",
    "                    big = alpha\n",
    "\n",
    "        # apply difference to potentials as needed\n",
    "        for dx in S:\n",
    "            rpotentials[dx] -= big\n",
    "\n",
    "        for dy in T:\n",
    "            cpotentials[dy] += big\n",
    "        \n",
    "    # Step 4: Loop while our matching is too small\n",
    "    while len(matching) != size:\n",
    "        # Step A: Compute neighbours in equality graph\n",
    "        NS = neighbours(S)\n",
    "        if set([b[1] for b in NS]) == T:\n",
    "            # Step B: If all firms are in the tree, update potentials to get a new one\n",
    "            update_potentials()\n",
    "            NS = neighbours(S)\n",
    "\n",
    "        # get the untreed firm\n",
    "        pair = next(n for n in NS if n[1] not in T)\n",
    "        if pair[1] not in [m[1] for m in matching]:\n",
    "            # Step D: Firm is not matched so add it to matching \n",
    "            matching.append(pair)\n",
    "            # Step E: Swap the alternating path in our alternating tree attached to the worker we matched\n",
    "            source = x_nodes[pair[0]]\n",
    "            matched = 1\n",
    "            while source.parent != None:\n",
    "                above = source.parent\n",
    "                if matched:\n",
    "                    # if previously matched, this should be removed from matching\n",
    "                    matching.remove([source.val, above.val])\n",
    "                else:\n",
    "                    # if previous was a remove, this is a match\n",
    "                    matching.append([above.val, source.val])\n",
    "\n",
    "                matched = 1 - matched\n",
    "                source = above\n",
    "\n",
    "            # Step F: Destroy the tree, go to Step 4 to check completion, and possibly go to Step A\n",
    "            free = list(set(range(size)) - set([m[0] for m in matching]))\n",
    "            if len(free):\n",
    "                tree_root = Node(free[0])\n",
    "                x_nodes = {free[0]: tree_root}\n",
    "                S = {free[0]}\n",
    "                T = set()\n",
    "\n",
    "        else:\n",
    "            # Step C: Firm is matched so add it to the tree and go back to Step A\n",
    "            matching_x = next(m[0] for m in matching if m[1] == pair[1])\n",
    "            S.add(matching_x)\n",
    "            T.add(pair[1])\n",
    "            source = x_nodes[pair[0]]\n",
    "            y_node = Node(pair[1], source)\n",
    "            x_node = Node(matching_x, y_node)\n",
    "            x_nodes[matching_x] = x_node\n",
    "    \n",
    "    revenues = [matrx[m[0], m[1]] for m in matching]\n",
    "    return (matching, revenues, rpotentials, cpotentials, sum(revenues))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the complete algorithm, which will work on any input so long as it has revenues for every worker-firm pair. Also note it requires the number of workers to be the same as the number of firms; if there are more workers or more firms, the matrix can be adjusted to be square by adding dummy firms or dummy workers that give revenues of zero.\n",
    "\n",
    "We can try running our algorithm on our simple example now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matching: [[0, 1], [1, 0]]\n",
      "Revenue: [50, 80]\n",
      "Wages: [50, 30]\n",
      "Profits: [50, 0]\n",
      "Revenue Sum: 130\n"
     ]
    }
   ],
   "source": [
    "match, rev, rpot, cpot, revsum = run_hungarian(np.array([[100, 50], [80, 26]]))\n",
    "print(\"Matching:\", match)\n",
    "print(\"Revenue:\", rev)\n",
    "print(\"Wages:\", rpot)\n",
    "print(\"Profits:\", cpot)\n",
    "print(\"Revenue Sum:\", revsum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is exactly the result from before, so the algorithm worked. We could try it on a bigger matrix, such as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matching: [[2, 5], [3, 3], [4, 7], [8, 1], [5, 2], [9, 0], [6, 6], [1, 4], [0, 9], [7, 8]]\n",
      "Revenue: [155, 169, 189, 176, 143, 193, 179, 166, 155, 180]\n",
      "Wages: [128, 166, 147, 169, 189, 143, 160, 147, 176, 139]\n",
      "Profits: [54, 0, 0, 0, 0, 8, 19, 0, 33, 27]\n",
      "Revenue Sum: 1705\n",
      "Took 0.0005710124969482422 seconds\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "matrx = np.random.randint(0, 200, (10, 10))\n",
    "a = time.time()\n",
    "match, rev, rpot, cpot, revsum = run_hungarian(matrx)\n",
    "b = time.time()\n",
    "print(\"Matching:\", match)\n",
    "print(\"Revenue:\", rev)\n",
    "print(\"Wages:\", rpot)\n",
    "print(\"Profits:\", cpot)\n",
    "print(\"Revenue Sum:\", revsum)\n",
    "print(\"Took\", b-a, \"seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which takes a random 10 by 10 matrix. To verify, one could brute-force the problem by manually checking every possible matching of workers with firms, which runs much slower. Verification was done using this method for smaller sizes, and by indexing the result against a popular existing Hungarian Algorithm library for larger sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.system(\"pip install hungarian-algorithm\")\n",
    "from hungarian_algorithm import algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to test their algorithm, we need to be able to convert our payoff matrix into a dictionary of dictionaries, which is the format they want for their input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "def convert(mat, size):\n",
    "    d = defaultdict(dict)\n",
    "    for i in range(size):\n",
    "        for j in range(size):\n",
    "            d[str(i)][str(j) + \"b\"] = mat[i, j]\n",
    "\n",
    "    return dict(d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can try it out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(('1', '4b'), 166), (('5', '2b'), 143), (('6', '6b'), 179), (('8', '1b'), 176), (('0', '9b'), 155), (('2', '5b'), 155), (('9', '0b'), 193), (('4', '7b'), 189), (('7', '8b'), 180), (('3', '3b'), 169)]\n",
      "[166, 143, 179, 176, 155, 155, 193, 189, 180, 169]\n",
      "[155, 169, 189, 176, 143, 193, 179, 166, 155, 180]\n",
      "1705\n",
      "1705\n",
      "Library took 0.017161130905151367 seconds\n",
      "Ours took 0.0005710124969482422 seconds\n"
     ]
    }
   ],
   "source": [
    "G = convert(matrx, matrx.shape[0])\n",
    "try:\n",
    "    c = time.time()\n",
    "    otherresult = algorithm.find_matching(G, matching_type = 'max', return_type = 'list')\n",
    "    d = time.time()\n",
    "    print(otherresult)\n",
    "    print([tup[1] for tup in otherresult])\n",
    "    print(rev)\n",
    "    #print(\"Sets equal?\", set([tup[1] for tup in otherresult]) == set(rev))\n",
    "    print(sum([tup[1] for tup in otherresult]))\n",
    "    print(sum(rev))\n",
    "except:\n",
    "    print(\"Library Crashed\")\n",
    "    \n",
    "\n",
    "print(\"Library took\", d-c, \"seconds\")\n",
    "print(\"Ours took\", b-a, \"seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In some random iterations, they found different matchings that have the same optimal result. When we have larger payoff matrices, it seems certainly likely that an optimal matching may not be unique if we have a few similar entries.\n",
    "\n",
    "Additionally, note the speed difference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a sanity check, lets run the two algorithms on a smaller matrix that is likely to not have this phenomenon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matching: [[0, 4], [2, 0], [3, 2], [1, 3], [4, 1]]\n",
      "Revenue: [195, 137, 199, 33, 143]\n",
      "Wages: [162, 33, 69, 168, 68]\n",
      "Profits: [68, 75, 31, 0, 33]\n",
      "Revenue Sum: 707\n",
      "Took 0.00035953521728515625 seconds\n"
     ]
    }
   ],
   "source": [
    "matrx = np.random.randint(0, 200, (5, 5))\n",
    "a = time.time()\n",
    "match, rev, rpot, cpot, revsum = run_hungarian(matrx)\n",
    "b = time.time()\n",
    "print(\"Matching:\", match)\n",
    "print(\"Revenue:\", rev)\n",
    "print(\"Wages:\", rpot)\n",
    "print(\"Profits:\", cpot)\n",
    "print(\"Revenue Sum:\", revsum)\n",
    "print(\"Took\", b-a, \"seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(('4', '1b'), 143), (('1', '3b'), 33), (('3', '2b'), 199), (('2', '0b'), 137), (('0', '4b'), 195)]\n",
      "[143, 33, 199, 137, 195]\n",
      "[195, 137, 199, 33, 143]\n",
      "707\n",
      "707\n",
      "Library took 0.004735469818115234 seconds\n",
      "Ours took 0.00035953521728515625 seconds\n"
     ]
    }
   ],
   "source": [
    "G = convert(matrx, matrx.shape[0])\n",
    "try:\n",
    "    c = time.time()\n",
    "    otherresult = algorithm.find_matching(G, matching_type = 'max', return_type = 'list')\n",
    "    d = time.time()\n",
    "    print(otherresult)\n",
    "    print([tup[1] for tup in otherresult])\n",
    "    print(rev)\n",
    "    #print(\"Sets equal?\", set([tup[1] for tup in otherresult]) == set(rev))\n",
    "    print(sum([tup[1] for tup in otherresult]))\n",
    "    print(sum(rev))\n",
    "except:\n",
    "    print(\"Library Crashed\")\n",
    "    \n",
    "\n",
    "print(\"Library took\", d-c, \"seconds\")\n",
    "print(\"Ours took\", b-a, \"seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, the results check out, or they did when tested before upload. If you clone this Jupyter notebook to your own machine, you can play around with the numbers, or pass in your own matrices, and compare the two algorithms yourself."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Special thanks to anonymous individuals for providing the following speed optimizations:*\n",
    "- consistent usage of set() types for S and T\n",
    "- faster NumPy random matrix generator\n",
    "- identification of redundant lines including copy and data from previous versions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Works Cited\n",
    "\n",
    "Golin, M. *Bipartite Matching & the Hungarian Method*. Personal collection of M. Golin, Hong Kong University of Science and Technology, Clear Water Bay, Kowloon.\n",
    " \n",
    "Peters, M. *Hungarian Algorithm*. Personal collection of M. Peters, University of British Columbia, Vancouver, BC."
   ]
  }
 ],
 "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}