Skip to content

Instantly share code, notes, and snippets.

@travishen

travishen/kruskal-python.ipynb

Last active December 13, 2018 07:25
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Embed
What would you like to do?
Learn more about clone URLs
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