【算法——Python实现】有权图求最小生成树Prim算法

class Edge(object):
"""边"""
def __init__(self, a, b, weight):
        self.a = a # 第一个顶点
        self.b = b # 第二个顶点
        self.weight = weight # 权值

def v(self):
return self.a

def w(self):
return self.b

def wt(self):
return self.weight

def other(self, x):
# 返回x顶点连接的另一个顶点
if x == self.a or x == self.b:
if x == self.a:
return self.b
else:
return self.a

def __lt__(self, other):
# 小于号重载
return self.weight < other.wt()

def __le__(self, other):
# 小于等于号重载
return self.weight <= other.wt()

def __gt__(self, other):
# 大于号重载
return self.weight > other.wt()

def __ge__(self, other):
# 大于等于号重载
return self.weight >= other.wt()

def __eq__(self, other):
# ==号重载
return self.weight == other.wt()


class DenseGraph(object):
"""有权稠密图 - 邻接矩阵"""
def __init__(self, n, directed):
        self.n = n  # 图中的点数
        self.m = 0  # 图中的边数
        self.directed = directed  # bool值，表示是否为有向图
        self.g = [[None for _ in range(n)] for _ in range(n)]  # 矩阵初始化都为None的二维矩阵

def V(self):
# 返回图中点数
return self.n

def E(self):
# 返回图中边数
return self.m

def addEdge(self, v, w, weight):
# v和w中增加一条边，v和w都是[0,n-1]区间
if v >= 0 and v < n and w >= 0 and w < n:
if self.hasEdge(v, w):
                self.m -= 1
            self.g[v][w] = Edge(v, w, weight)
if not self.directed:
                self.g[w][v] = Edge(w, v, weight)
            self.m += 1

def hasEdge(self, v, w):
# v和w之间是否有边，v和w都是[0,n-1]区间
if v >= 0 and v < n and w >= 0 and w < n:
return self.g[v][w] != None

class adjIterator(object):
"""相邻节点迭代器"""
def __init__(self, graph, v):
            self.G = graph  # 需要遍历的图
            self.v = v  # 遍历v节点相邻的边
            self.index = 0  # 遍历的索引

def __iter__(self):
return self

def next(self):
while self.index < self.G.V():
# 当索引小于节点数量时遍历，否则为遍历完成，停止迭代
if self.G.g[self.v][self.index]:
                    r = self.G.g[self.v][self.index]
                    self.index += 1
return r
                self.index += 1
raise StopIteration()


class SparseGraph(object):
"""有权稀疏图- 邻接表"""
def __init__(self, n, directed):
        self.n = n  # 图中的点数
        self.m = 0  # 图中的边数
        self.directed = directed  # bool值，表示是否为有向图
        self.g = [[] for _ in range(n)]  # 矩阵初始化都为空的二维矩阵

def V(self):
# 返回图中点数
return self.n

def E(self):
# 返回图中边数
return self.m

def addEdge(self, v, w, weight):
# v和w中增加一条边，v和w都是[0,n-1]区间
if v >= 0 and v < n and w >= 0 and w < n:
# 考虑到平行边会让时间复杂度变为最差为O(n)
# if self.hasEdge(v, w):
# return None
            self.g[v].append(Edge(v, w, weight))
if v != w and not self.directed:
                self.g[w].append(Edge(w, v, weight))
            self.m += 1

def hasEdge(self, v, w):
# v和w之间是否有边，v和w都是[0,n-1]区间
# 时间复杂度最差为O(n)
if v >= 0 and v < n and w >= 0 and w < n:
for i in self.g[v]:
if i.other(v) == w:
return True
return False

class adjIterator(object):
"""相邻节点迭代器"""
def __init__(self, graph, v):
            self.G = graph  # 需要遍历的图
            self.v = v  # 遍历v节点相邻的边
            self.index = 0  # 遍历的索引

def __iter__(self):
return self

def next(self):
if len(self.G.g[self.v]):
# v有相邻节点才遍历
if self.index < len(self.G.g[self.v]):
                    r = self.G.g[self.v][self.index]
                    self.index += 1
return r
else:
raise StopIteration()
else:
raise StopIteration()


class ReadGraph(object):
"""读取文件中的图"""
def __init__(self, graph, filename):
with open(filename, 'r') as f:
            line = f.readline()
            line = line.strip('\n')
            line = line.split()
            v = int(line[0])
            e = int(line[1])
if v == graph.V():
                lines = f.readlines()
for i in lines:
                    a, b, w = self.stringstream(i)
if a >= 0 and a < v and b >=0 and b < v:
                        graph.addEdge(a, b, w)

def stringstream(self, text):
        result = text.strip('\n')
        result = result.split()
        a, b, w = result
return int(a), int(b), float(w)


class IndexMinHeap(object):
"""
 最小反向索引堆，出堆不删除data，只删除indexs
 """
def __init__(self, n):
        self.capacity = n # 堆的最大容量
        self.data = [-1 for _ in range(n)]  # 创建堆
        self.indexs = [] # 创建索引堆
        self.reverse = [-1 for _ in range(n)] # 创建反向索引
        self.count = 0  # 元素数量

def size(self):
return self.count

def isEmpty(self):
return self.count == 0

def insert(self, i, item):
# 插入元素入堆
        self.data[i] = item
        self.indexs.append(i)
        self.reverse[i] = self.count
        self.count += 1
        self.shiftup(self.count)

def shiftup(self, count):
# 将插入的元素放到合适位置，保持最小堆
while count > 1 and self.data[self.indexs[(count/2)-1]] > self.data[self.indexs[count-1]]:
            self.indexs[(count/2)-1], self.indexs[count-1] = self.indexs[count-1], self.indexs[(count/2)-1]
            self.reverse[self.indexs[(count/2)-1]] = (count/2)-1
            self.reverse[self.indexs[count-1]] = count-1
            count /= 2

def extractMin(self):
# 出堆
if self.count > 0:
            ret = self.data[self.indexs[0]]
            self.indexs[0], self.indexs[self.count-1] = self.indexs[self.count-1], self.indexs[0]
            self.reverse[self.indexs[0]] = 0
            self.reverse[self.indexs[self.count-1]] = self.count-1
# for i in range(self.indexs[self.count-1]+1, self.count):
# self.indexs[self.reverse[i]] -= 1
            self.reverse[self.indexs[self.count-1]] = -1
# self.data[self.indexs[self.count-1]] = -1
            self.indexs.pop(self.count-1)
            self.count -= 1
            self.shiftDown(1)
return ret

def extractMinIndex(self):
# 出堆返回索引
if self.count > 0:
            ret = self.indexs[0]
            self.indexs[0], self.indexs[self.count-1] = self.indexs[self.count-1], self.indexs[0]
            self.reverse[self.indexs[0]] = 0
            self.reverse[self.indexs[self.count-1]] = self.count-1
# for i in range(self.indexs[self.count-1]+1, self.count):
# self.indexs[self.reverse[i]] -= 1
            self.reverse[self.indexs[self.count-1]] = -1
# self.data[self.indexs[self.count-1]] = -1
            self.indexs.pop(self.count-1)
            self.count -= 1
            self.shiftDown(1)
return ret

def shiftDown(self, count):
# 将堆的索引位置元素向下移动到合适位置，保持最小堆
while 2 * count <= self.count :
# 证明有孩子
            j = 2 * count
if j + 1 <= self.count:
# 证明有右孩子
if self.data[self.indexs[j]] < self.data[self.indexs[j-1]]:
# 右孩子数值比左孩子数值小
                    j += 1
if self.data[self.indexs[count-1]] <= self.data[self.indexs[j-1]]:
# 堆的索引位置已经小于两个孩子节点，不需要交换了
break
            self.indexs[count-1], self.indexs[j-1] = self.indexs[j-1], self.indexs[count-1]
            self.reverse[self.indexs[count-1]] = count-1
            self.reverse[self.indexs[j-1]] =j-1
            count = j

def getItem(self, i):
# 根据索引获取数值
if i >=0 and i <= self.count-1:
return self.data[i]
else:
return None

def change(self, i, newItem):
# 改变i索引位置的数值
if i >=0:
            self.data[i] = newItem
            j = self.reverse[i]
            self.shiftup(j+1)
            self.shiftDown(j+1)


class PrimMST(object):
"""最小生成树"""
def __init__(self, graph):
        self.G = graph # 传入图
        self.ipq = IndexMinHeap(self.G.V()) # 最小索引堆，存权值
        self.edgeTo = [None for _ in range(self.G.V())] # 存Edge边，表示存放连接索引节点的边
        self.marked = [False for _ in range(self.G.V())]  # 用于标记已经被选取为树的节点，初始都为False
        self.mst = []  # 记录被选取的边
        self.mstWeight = 0  # 记录最小生成树的总权值

# Prim
        self.visit(0)
while not self.ipq.isEmpty():
# 取出权值最小的横切边的索引，该索引表示这条边与哪个节点相连接
            v = self.ipq.extractMinIndex()
            self.mst.append(self.edgeTo[v])
# 即将将节点标记为已访问（红色）
            self.visit(v)
        self.mstWeight = sum([i.wt() for i in self.mst])

def visit(self, v):
# 访问
if not self.marked[v]:
# self.marked[v]为False，表示该节点还未被加入树中
            self.marked[v] = True
            adj = self.G.adjIterator(self.G, v)
for i in adj:
# 与v相连的另一端
                w = i.other(v)
if not self.marked[w]:
# 与v相连接的另一端还未被加入树中，则这一条边为横切边
if not self.edgeTo[w]:
# 如果连接w端点的边还没有横切边，则直接把权值插入到最小索引堆
                        self.ipq.insert(w, i.wt())
                        self.edgeTo[w] = i
elif i.wt() < self.edgeTo[w].wt():
# 如果有连接w端点的横切边，且新的横切边的权值小于原有横切边的权值，则用新的横切边替代
                        self.edgeTo[w] = i
                        self.ipq.change(w, i.wt())

def mstEdges(self):
# 查询最小生成树的边
return self.mst

def result(self):
# 返回最小生成树的权值
return self.mstWeight