LazyPrimMST.spanning_tree()   A
last analyzed

Complexity

Conditions 1

Size

Total Lines 2

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 1
c 1
b 0
f 0
dl 0
loc 2
rs 10
1
from abc import ABCMeta, abstractmethod
2
3
from pyalgs.algorithms.commons.union_find import UnionFind
4
from pyalgs.algorithms.commons.util import less
5
from pyalgs.data_structures.commons.bag import Bag
6
from pyalgs.data_structures.commons.priority_queue import MinPQ, IndexMinPQ
7
from pyalgs.data_structures.graphs.graph import EdgeWeightedGraph
8
9
10
class MST(object):
11
    __metaclass__ = ABCMeta
12
13
    @abstractmethod
14
    def spanning_tree(self):
15
        pass
16
17
18
class KruskalMST(MST):
19
    tree = None
20
21
    def __init__(self, G):
22
        if not isinstance(G, EdgeWeightedGraph):
23
            raise ValueError('Graph must be edge weighted and undirected to run MST')
24
        minpq = MinPQ.create()
25
        self.tree = Bag()
26
        for e in G.edges():
27
            minpq.enqueue(e)
28
29
        uf = UnionFind.create(G.vertex_count())
30
31
        while not minpq.is_empty() and self.tree.size() < G.vertex_count() - 1:
32
            e = minpq.del_min()
33
            v = e.either()
34
            w = e.other(v)
35
            if not uf.connected(v, w):
36
                uf.union(v, w)
37
                self.tree.add(e)
38
39
    def spanning_tree(self):
40
        return self.tree.iterate()
41
42
43
class LazyPrimMST(MST):
44
    tree = None
45
    marked = None
46
    minpq = None
47
48
    def __init__(self, G):
49
        if not isinstance(G, EdgeWeightedGraph):
50
            raise ValueError('Graph must be edge weighted and undirected to run MST')
51
        self.minpq = MinPQ.create()
52
        self.tree = Bag()
53
        vertex_count = G.vertex_count()
54
        self.marked = [False] * vertex_count
55
        self.visit(G, 0)
56
57
        while not self.minpq.is_empty() and self.tree.size() < vertex_count - 1:
58
            edge = self.minpq.del_min()
59
            v = edge.either()
60
            w = edge.other(v)
61
            if self.marked[v] and self.marked[w]:
62
                continue
63
            self.tree.add(edge)
64
            if not self.marked[v]:
65
                self.visit(G, v)
66
            if not self.marked[w]:
67
                self.visit(G, w)
68
69
    def visit(self, G, v):
70
        self.marked[v] = True
71
        for e in G.adj(v):
72
            w = e.other(v)
73
            if not self.marked[w]:
74
                self.minpq.enqueue(e)
75
76
    def spanning_tree(self):
77
        return self.tree.iterate()
78
79
80
class EagerPrimMST(MST):
81
    path = None
82
    pq = None
83
    marked = None
84
85
    def __init__(self, G):
86
        if not isinstance(G, EdgeWeightedGraph):
87
            raise ValueError('Graph must be edge weighted and undirected to run MST')
88
        vertex_count = G.vertex_count()
89
        self.pq = IndexMinPQ(vertex_count)
90
        self.path = Bag()
91
        self.marked = [False] * vertex_count
92
93
        self.visit(G, 0)
94
95
        while not self.pq.is_empty() and self.path.size() < vertex_count - 1:
96
            e = self.pq.min_key()
97
            w = self.pq.del_min()
98
            self.path.add(e)
99
            if not self.marked[w]:
100
                self.visit(G, w)
101
102
    def visit(self, G, v):
103
        self.marked[v] = True
104
        for e in G.adj(v):
105
            w = e.other(v)
106
            if not self.marked[w]:
107
                if self.pq.contains_index(w):
108
                    old_e = self.pq.get(w)
109
                    if less(e, old_e):
110
                        self.pq.decrease_key(w, e)
111
                else:
112
                    self.pq.insert(w, e)
113
114
    def spanning_tree(self):
115
        return self.path.iterate()
116