union_find.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Union-Find Data Structures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper Functions for Experiments and Drawing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use again `networkx` for representing and drawing graphs/trees. In addition, we use `pygraphviz` for a better plotting of trees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from networkx.drawing.nx_agraph import graphviz_layout\n",
    "import networkx as nx\n",
    "import pygraphviz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We randomly create a small example input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "import random\n",
    "\n",
    "class Example:\n",
    "    def __init__(self, no_nodes, no_union_calls):\n",
    "        self.no_nodes = no_nodes\n",
    "        self.union_calls = []\n",
    "        \n",
    "        node_pairs = list(itertools.combinations(range(no_nodes), 2))\n",
    "        sampled_pairs = random.sample(node_pairs, no_union_calls)\n",
    "        \n",
    "        # In the sampled pairs, the node with the smaller number\n",
    "        # always is the first component. We thus swap the order\n",
    "        # for some of the pairs       \n",
    "        for node1, node2 in sampled_pairs:\n",
    "            if random.choice((True, False)):\n",
    "                node1, node2 = node2, node1\n",
    "            self.union_calls.append((node1, node2))\n",
    "\n",
    "small_example = Example(10, 8)\n",
    "print(small_example.union_calls)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To be able to quickly verify the results of the algorithms below, we visualize the resulting graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = nx.Graph()\n",
    "graph.add_nodes_from(range(small_example.no_nodes))\n",
    "graph.add_edges_from(small_example.union_calls)\n",
    "pos = nx.spring_layout(graph, k=1)\n",
    "# if you have scipy installed, you can also use:\n",
    "# pos = nx.kamada_kawai_layout(graph)\n",
    "nx.draw(graph, pos, with_labels=True, \n",
    "        node_size=1000, node_color='lightblue')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Quick-Find"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Quick-Find is a rather naive approach that stores for each vertex the representative of the corresponding connected component. If two connected components are merged, we need to iterate over all vertices and update for one of the two components the representative to the representative of the other component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class QuickFind:\n",
    "    def __init__(self, no_nodes):\n",
    "        self.representative = list(range(no_nodes))\n",
    "        self.components = no_nodes\n",
    "\n",
    "    def find(self, v):\n",
    "        return self.representative[v]\n",
    "\n",
    "    def union(self, v, w):\n",
    "        repr_v = self.find(v)\n",
    "        repr_w = self.find(w)\n",
    "        if repr_v == repr_w:  # already in same component\n",
    "            return\n",
    "        # replace all occurrences of repr_v in self.representative with repr_w\n",
    "        for i in range(len(self.representative)):\n",
    "            if self.representative[i] == repr_v:\n",
    "                self.representative[i] = repr_w\n",
    "        self.components -= 1  # we merged two components\n",
    "        \n",
    "    def connected(self, v, w):\n",
    "        return self.find(v) == self.find(w)\n",
    "\n",
    "    def count(self):\n",
    "        return self.components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qf = QuickFind(small_example.no_nodes)\n",
    "print(\"Initial representatives:\", qf.representative)\n",
    "print()\n",
    "\n",
    "for x, y in small_example.union_calls:\n",
    "    qf.union(x, y)\n",
    "    print(\"after union(%i, %i):\" % (x, y))\n",
    "    print(qf.representative)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "qf.connected(0, 6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quick-Union"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Quick-Union stores all vertices of a connected component in a tree. The advantage is that in `union` one does not have to traverse all vertices. In `find`, we can no longer directly access the representative, but use instead the root of the corresponding tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class QuickUnion:\n",
    "    def __init__(self, no_nodes):\n",
    "        self.parent = list(range(no_nodes))\n",
    "        self.components = no_nodes\n",
    "\n",
    "    def find(self, v):\n",
    "        while self.parent[v] != v:\n",
    "            v = self.parent[v]\n",
    "        return v\n",
    "\n",
    "    def union(self, v, w):\n",
    "        repr_v = self.find(v)\n",
    "        repr_w = self.find(w)\n",
    "        if repr_v == repr_w:  # already in same component\n",
    "            print(\"Vertices\", v, \"and\", w, \"already were in\",\n",
    "                  \"the same connected component.\")\n",
    "            return\n",
    "        self.parent[repr_v] = repr_w\n",
    "        print(\"Make the root\", repr_v, \"of the tree of\", v, \n",
    "              \"a child of the root\", repr_w, \"of the tree of\", w)\n",
    "        self.components -= 1\n",
    "        \n",
    "    def connected(self, v, w):\n",
    "        return self.find(v) == self.find(w)\n",
    "\n",
    "    def count(self):\n",
    "        return self.components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first write a small helper function for drawing the forest represented by `parent`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def draw_forest(parent_array):\n",
    "    graph = nx.DiGraph()\n",
    "    graph.add_nodes_from(range(len(parent_array)))\n",
    "    for child, parent in enumerate(parent_array):\n",
    "        if child != parent:\n",
    "            graph.add_edge(child, parent)\n",
    "        \n",
    "    pos = graphviz_layout(graph.reverse(), prog='dot')\n",
    "    nx.draw(graph, pos, with_labels=True, node_size=300, node_color='lightblue')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now test Quick-Union:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qu = QuickUnion(small_example.no_nodes)\n",
    "for x, y in small_example.union_calls:\n",
    "    print(\"union(%i, %i)\" % (x, y))\n",
    "    qu.union(x, y)\n",
    "    draw_forest(qu.parent)\n",
    "    plt.show()\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ranked Quick-Union"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class RankedQuickUnion:\n",
    "    def __init__(self, no_nodes):\n",
    "        self.parent = list(range(no_nodes))\n",
    "        self.components = no_nodes\n",
    "        self.rank = [0] * no_nodes  # [0, ..., 0]\n",
    "\n",
    "    def find(self, v):\n",
    "        while self.parent[v] != v:\n",
    "            v = self.parent[v]\n",
    "        return v\n",
    "\n",
    "    def union(self, v, w):\n",
    "        repr_v = self.find(v)\n",
    "        repr_w = self.find(w)\n",
    "        if repr_v == repr_w:  # already in same component\n",
    "            print(\"Vertices\", v, \"and\", w, \"already were in\",\n",
    "                  \"the same connected component.\")\n",
    "            return\n",
    "\n",
    "        if self.rank[repr_w] < self.rank[repr_v]:\n",
    "            self.parent[repr_w] = repr_v\n",
    "            print(\"Make the root\", repr_w, \"of the tree of\", w, \n",
    "                  \"a child of the root\", repr_v, \"of the tree of\", v)\n",
    "        else:\n",
    "            self.parent[repr_v] = repr_w\n",
    "            print(\"Make the root\", repr_v, \"of the tree of\", v, \n",
    "                  \"a child of the root\", repr_w, \"of the tree of\", w)\n",
    "\n",
    "            if self.rank[repr_v] == self.rank[repr_w]:\n",
    "                print(\"Update the height of the resulting tree\")\n",
    "                self.rank[repr_w] += 1\n",
    "        self.components -= 1\n",
    "\n",
    "    def connected(self, v, w):\n",
    "        return self.find(v) == self.find(w)\n",
    "\n",
    "    def count(self):\n",
    "        return self.components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rqu = RankedQuickUnion(small_example.no_nodes)\n",
    "for x, y in small_example.union_calls:\n",
    "    print(\"union(%i, %i)\" % (x, y))\n",
    "    rqu.union(x, y)\n",
    "    draw_forest(rqu.parent)\n",
    "    plt.show()\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ranked Quick-Union with Path Compression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Path compression is an improvement that makes the trees more shallow as a side effect of a `find` call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class RankedQuickUnionWithPathCompression:\n",
    "    def __init__(self, no_nodes):\n",
    "        self.parent = list(range(no_nodes))\n",
    "        self.components = no_nodes\n",
    "        self.rank = [0] * no_nodes  # [0, ..., 0]\n",
    "\n",
    "    def find(self, v):\n",
    "        if self.parent[v] == v:\n",
    "            return v\n",
    "        root = self.find(self.parent[v])\n",
    "        self.parent[v] = root\n",
    "        return root\n",
    "\n",
    "    def union(self, v, w):\n",
    "        repr_v = self.find(v)\n",
    "        repr_w = self.find(w)\n",
    "        if repr_v == repr_w:\n",
    "            return\n",
    "        if self.rank[repr_w] < self.rank[repr_v]:\n",
    "            self.parent[repr_w] = repr_v\n",
    "        else:\n",
    "            self.parent[repr_v] = repr_w\n",
    "            if self.rank[repr_v] == self.rank[repr_w]:\n",
    "                self.rank[repr_w] += 1\n",
    "        self.components -= 1\n",
    "\n",
    "    def connected(self, v, w):\n",
    "        return self.find(v) == self.find(w)\n",
    "\n",
    "    def count(self):\n",
    "        return self.components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Method `union` is implemented as without path compression. But since `union` internally calls `find`, it also reconnects vertices to the root."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rqupc = RankedQuickUnionWithPathCompression(small_example.no_nodes)\n",
    "for x, y in small_example.union_calls:\n",
    "    rqupc.union(x, y)\n",
    "    \n",
    "draw_forest(rqupc.parent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Without path compression, the calls to `find` do not modify the data. With path compression, the trees get more and more shallow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# EXERCISE identify for the randomly generated example suitable calls to find and connected:\n",
    "rqupc.connected(1, 7)\n",
    "rqupc.find(5)\n",
    "draw_forest(rqupc.parent)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}