Source code for tryalgo.kruskal

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""\
Minimum spanning tree by kruskal
jill-jenn vie et christoph durr - 2014-2018
"""

from math import sqrt
import random


# snip{ union-find
[docs]class UnionFind: """Maintains a partition of {0, ..., n-1} """ def __init__(self, n): self.up_bound = list(range(n)) self.rank = [0] * n
[docs] def find(self, x_index): """ :returns: identifier of part containing x_index :complex_indexity: O(inverse_ackerman(n)) """ if self.up_bound[x_index] == x_index: return x_index self.up_bound[x_index] = self.find(self.up_bound[x_index]) return self.up_bound[x_index]
[docs] def union(self, x_index, y_index): """ Merges part that contain x and part containing y :returns: False if x_index, y_index are already in same part :complexity: O(inverse_ackerman(n)) """ repr_x = self.find(x_index) repr_y = self.find(y_index) if repr_x == repr_y: # already in the same component return False if self.rank[repr_x] == self.rank[repr_y]: self.rank[repr_x] += 1 self.up_bound[repr_y] = repr_x elif self.rank[repr_x] > self.rank[repr_y]: self.up_bound[repr_y] = repr_x else: self.up_bound[repr_x] = repr_y return True
# snip} # snip{ kruskal # pylint: disable=redefined-outer-name, unused-variable
[docs]def kruskal(graph, weight): """Minimum spanning tree by Kruskal :param graph: undirected graph in listlist or listdict format :param weight: in matrix format or same listdict graph :returns: list of edges of the tree :complexity: ``O(|E|log|E|)`` """ u_f = UnionFind(len(graph)) edges = [] for u, _ in enumerate(graph): for v in graph[u]: edges.append((weight[u][v], u, v)) edges.sort() min_span_tree = [] for w_idx, u_idx, v_idx in edges: if u_f.union(u_idx, v_idx): min_span_tree.append((u_idx, v_idx)) return min_span_tree
# snip}
[docs]def dist(a, b): """ distance between point a and point b """ return sqrt(sum([(a[i] - b[i]) * (a[i] - b[i]) for i in range(len(a))]))
# pylint: disable=pointless-string-statement if __name__ == "__main__": """ main function """ N = 256 points = [[random.random() * 5, random.random() * 5] for _ in range(N)] weight = [[dist(points[i], points[j]) for j in range(N)] for i in range(N)] graph = [[j for j in range(N) if i != j] for i in range(N)] with open('../data/kruskal-points.tex', 'w') as infile: min_span_tree = kruskal(graph, weight) val = 0 for u_idx, v_idx in min_span_tree: val += weight[u_idx][v_idx] infile.write('\\draw[blue] (%f, %f) -- (%f, %f);\n' % tuple(points[u_idx] + points[v_idx])) for point in points: infile.write('\\filldraw[black] (%f, %f) circle (1pt);\n' % tuple(point)) print(val)