python 使用 最大生成树 解决 营救问题

题目描述

你是红军指挥官，在一场军事演习中，你的部分军队被蓝军包围了。蓝军包围的方式如下

在上图中，每个顶点表示蓝军的部队，顶点中数字表示蓝军在此处的人数（千人），两点间的边表示蓝军两个部队间形成的火线，火线构成的圈即是一道包围，一条火线的战斗力为其相连两个部队的人数和，也是你要进攻这条火线所要消耗的兵力。你可以同时进攻蓝军的多条火线，请以成本最低的方式打破蓝军包围，营救被包围的部队，计算出所需要消耗的兵力数（千人）

输入

• 输入包含多个测试例。第一行是一个整数n，表示测试测试例数量。
• 对每个测试例：
  第一行是一个正整数m(2<m≤200),是蓝军部队数。
  每个蓝军部队有两行：
  一行3个正整数:i为部队编号（0≤i≤m-1）,ai为蓝军部队人数（千人）(1≤ai≤100)，bi为与该部队能形成火线的部队数量(1≤bi≤m-1)。
  一行有bi个正整数，分别是与部队i形成火线的部队编号，数字间有空格。
• 至少有一个包围。

输出

每个测试例输出一行，是一个正整数：打破包围营救战友的最小消耗兵力。

输入样例

1 
3  
0 1 2 
1 2 
1 2 2 
0 2 
2 3 2 
0 1

输出结果

3

解题思路

代码实现

使用并查集优化的Kruskal算法实现

#!/usr/bin/env python3
#rescue.py
#fumiama 20201115

from heapq import heappop, heappush

class UnionFindSet(object):
    def __init__(self, nodes):
        self.fatherMap = {}
        self.setNumMap = {}
        for node in nodes:
            self.fatherMap[node] = node
            self.setNumMap[node] = 1
    def findFather(self, node):
        father = self.fatherMap[node]
        if (node != father):
            father = self.findFather(father)
        self.fatherMap[node] = father
        return father
    def isSameSet(self, a, b):
        return self.findFather(a) == self.findFather(b)
    def union(self, a, b):
        if a is None or b is None: return
        aFather = self.findFather(a)
        bFather = self.findFather(b)
        if (aFather != bFather):
            aNum = self.setNumMap[aFather]
            bNum = self.setNumMap[bFather]
            if (aNum <= bNum):
                self.fatherMap[aFather] = bFather
                self.setNumMap[bFather] = aNum + bNum
                self.setNumMap.pop(aFather)
            else:
                self.fatherMap[bFather] = aFather
                self.setNumMap[aFather] = aNum + bNum
                self.setNumMap.pop(bFather)

class Node(object):
    def __init__(self, name, size, links):
        self.name = name
        self.size = size
        self.links = links
    def __repr__(self):
        return "(节点:" + str(self.name) + ", 大小:" + str(self.size) + ", 连接到:" + str(self.links) + ")"

class Graph(object):
    def __init__(self):
        self.edges = []
        self.nodes = set()
        self.sizes = dict(map=int)
        self._edgesSets = []        #识别重复边
    def addNode(self, node):
        self.nodes.add(node)
        self.sizes[node.name] = node.size
    def calcEdgeWeights(self):
        for node in self.nodes:
            for link in node.links:
                edge = (-(node.size + self.sizes[link]), node.name, link)   #边权为负构造大顶堆
                sedge = {node.size + self.sizes[link], node.name, link}
                if sedge not in self._edgesSets:
                    heappush(self.edges, edge)
                    self._edgesSets.append(sedge)
    def rescue(self):
        forest = UnionFindSet(self.sizes.keys())
        edges = self.edges.copy()
        cost = 0
        while edges:
            edge = heappop(edges)
            if forest.isSameSet(edge[1], edge[2]): cost -= edge[0]
            else: forest.union(edge[1], edge[2])
        return cost

    def __repr__(self):
        return "图信息:\n边: " + str(self.edges) + "\n点: " + str(self.nodes)


if __name__ == '__main__':
    n = int(input())
    graph = Graph()
    while n:
        m = int(input())
        while m:
            no, size, edgeCnt = map(int, input().split())
            links = [int(x) for x in input().split()]
            graph.addNode(Node(no, size, links))
            m -= 1
        graph.calcEdgeWeights()
        #print(graph)
        minCost = graph.rescue()
        #print("最小兵力:", minCost)
        print(minCost)
        n -= 1