Skip to content

Instantly share code, notes, and snippets.

@travishen

travishen/kruskal-python.ipynb

Last active December 13, 2018 07:25
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save travishen/aa60d08a668691b83244caa395ddec07 to your computer and use it in GitHub Desktop.
Save travishen/aa60d08a668691b83244caa395ddec07 to your computer and use it in GitHub Desktop.
Download ZIP
Minimum Spinning Tree Implement using Kruskal by Python
Display the source blob
Display the rendered blob
Raw
kruskal-python.ipynb
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"class Edge:\n",
" \"\"\"Sortable edge in the graph\"\"\"\n",
" def __init__(self, weight, start, target):\n",
" self.weight = weight \n",
" self.start = start # Node\n",
" self.target = target # Node\n",
" \n",
" def __repr__(self):\n",
" return 'Edge(weight={}, start={}, target={})'.format(self.weight,\n",
" self.start,\n",
" self.target)\n",
" \n",
" def __cmp__(self, other):\n",
" return self.cmp(self.weight, other.weight)\n",
" \n",
" def __lt__(self, other):\n",
" return self.weight < other.weight\n",
" \n",
"class Node:\n",
" \"\"\"Node live in a graph / disjoint set\"\"\"\n",
" def __init__(self, name):\n",
" self.name = name\n",
" self.parent = None\n",
" self.set_ = None\n",
" \n",
" def __repr__(self):\n",
" return self.name\n",
" parent = None\n",
" if self.parent:\n",
" parent = self.parent.name\n",
" return 'Node(name={}, parent={})'.format(self.name, parent)\n",
" \n",
"class DisjointSet:\n",
" \"\"\"Represent a disjoint set\"\"\"\n",
" def __init__(self, node):\n",
" \"\"\"make set\"\"\"\n",
" self.nodes = set([node])\n",
" self.root = node \n",
" self.root.set_ = self\n",
" \n",
" def __str__(self):\n",
" if not self.nodes:\n",
" return 'Empty'\n",
" return str(self.nodes)\n",
" \n",
" def __len__(self):\n",
" return len(self.nodes)\n",
" \n",
" @staticmethod\n",
" def find(node):\n",
" \"\"\"Find root node in nodes and do path compression\"\"\"\n",
" \n",
" root = node\n",
" while root.parent is not None:\n",
" root = root.parent\n",
" \n",
" # path compression\n",
" while node is not root:\n",
" temp = node.parent\n",
" node.parent = root\n",
" node = temp\n",
" \n",
" return root\n",
" \n",
" @staticmethod\n",
" def merge(s1, s2):\n",
" \"\"\"Merge two set base on \"\"\"\n",
" \n",
" if s1 is s2: # is equal\n",
" return\n",
" \n",
" if len(s1) < len(s2): # s1 --> s2\n",
" s1.root.parent = s2.root \n",
" \n",
" for n in s1.nodes: # point all node to new set\n",
" n.set_ = s2 \n",
" s2.nodes.update(s1.nodes)\n",
" s1.nodes = set()\n",
" \n",
" else: # s2 --> s1\n",
" s2.root.parent = s1.root\n",
" \n",
" for n in s2.nodes: # point all node to new set\n",
" n.set_ = s1 \n",
" s1.nodes.update(s2.nodes)\n",
" s2.nodes = set()"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"class Kruskal:\n",
" def __init__(self, nodes, edges):\n",
" self.spanning_tree = []\n",
" self.edges = edges\n",
" self.edges.sort() # O(NlogN)\n",
" self.sets = []\n",
" self.nodes = nodes\n",
" for node in nodes:\n",
" self.sets.append(DisjointSet(node))\n",
" \n",
" def run(self):\n",
" \n",
" self.logging(init=True)\n",
" \n",
" for edge in self.edges:\n",
" r1 = DisjointSet.find(edge.start)\n",
" r2 = DisjointSet.find(edge.target)\n",
" \n",
" if r1.set_ is not r2.set_: # if in different set\n",
" \n",
" self.logging()\n",
"\n",
" DisjointSet.merge(r1.set_, r2.set_)\n",
" self.spanning_tree.append(edge)\n",
"\n",
" if len(self.spanning_tree) == len(self.nodes)-1: \n",
" # If we have selected n-1 edges, all the other \n",
" # edges will be discarted, so, we can stop here\n",
" self.logging()\n",
" break\n",
" \n",
" def logging(self, init=False):\n",
" row_format = ''\n",
" length = 0\n",
" for i, set_ in enumerate(self.sets):\n",
" if i == 2:\n",
" row_format += '{!s:>25} '\n",
" length += 26\n",
" else:\n",
" row_format += '{!s:>7} '\n",
" length += 8\n",
" \n",
" if init:\n",
" print(row_format.format(*range(0, len(self.sets))))\n",
" print('=' * length)\n",
" else:\n",
" print(row_format.format(*self.sets))"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"graph = []\n",
"\n",
"# construct A,B,C,D,E,F,G Nodes\n",
"node_str = 'ABCDEFG'\n",
"for s in node_str:\n",
" node = Node(s)\n",
" locals()[s] = node\n",
" graph.append(node)\n",
" \n",
"# lined nodes\n",
"edges = [\n",
" Edge(2, A, B),\n",
" Edge(6, A, C),\n",
" Edge(5, A, E),\n",
" Edge(10, A, F),\n",
" Edge(3, B, D),\n",
" Edge(3, B, E),\n",
" Edge(1, C, D),\n",
" Edge(2, C, F),\n",
" Edge(4, D, E),\n",
" Edge(5, D, G),\n",
" Edge(5, F, G),\n",
"]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"![image](https://storage.googleapis.com/ssivart/super9-blog/kruskal.png)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"algorithm = Kruskal(graph, edges)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## sets merging history"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
" 0 1 2 3 4 5 6 \n",
"==========================================================================\n",
" {A} {B} {C} {D} {E} {F} {G} \n",
" {A} {B} {D, C} Empty {E} {F} {G} \n",
" {A, B} Empty {D, C} Empty {E} {F} {G} \n",
" {A, B} Empty {F, D, C} Empty {E} Empty {G} \n",
" Empty Empty {F, D, A, B, C} Empty {E} Empty {G} \n",
" Empty Empty {F, A, B, C, D, E} Empty Empty Empty {G} \n",
" Empty Empty {F, A, G, B, C, D, E} Empty Empty Empty Empty \n"
]
}
],
"source": [
"algorithm.run()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## spinning tree edges"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"C ---> D\n",
"A ---> B\n",
"C ---> F\n",
"B ---> D\n",
"B ---> E\n",
"D ---> G\n"
]
}
],
"source": [
"for edge in algorithm.spanning_tree:\n",
" print('{} ---> {}'.format(edge.start, edge.target))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## set tree structure"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"A ---> C\n",
"B ---> C\n",
"C ---> None\n",
"D ---> C\n",
"E ---> C\n",
"F ---> C\n",
"G ---> C\n"
]
}
],
"source": [
"for node in graph:\n",
" parent = None\n",
" if node.parent:\n",
" parent = node.parent.name\n",
" print('{} ---> {}'.format(node.name, parent))"
]
}
],
"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.5.2"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment