# Classe MaxTree

Aug 17th, 2022
765
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. # file name: prim.py
2.
3. from typing import Dict, List
4.
5. # Class for a Directed Maximum weighted tree
6. class MaxTree:
7.   def __init__(self, nodes:List, costs:Dict):
8.     self.q = nodes.copy()   # List of all Nodes, will be empty at the end
9.     self.edges_cost = costs # Dict such as {(a, b): n}, where a, b are nodes and n is the weight between them
10.     self.tree_nodes = []    # List of all nodes in the tree
11.     self.tree_edges = []    # List of edges(tuples of nodes) on the tree
12.     self.children_of = {}   # Dict of children for each node
13.     self.tree_root = None   # Root node of the tree
14.
15.     self.__build_tree()
16.
17.   ################################
18.   # Using Prim's algorithm to create a maximum weighted spanning tree
19.   def __build_tree(self):
20.     highest_connection = {} # E
21.     cost_connection = {}    # C
22.
23.     for node in self.q:
24.       cost_connection[node] = -1
25.       highest_connection[node] = None
26.
27.     ################################
28.     # Applying the algorithm
29.     while len(self.q) > 0:
30.       node = max(self.q, key=lambda node: cost_connection[node])
31.       self.q.remove(node)
32.
33.       self.tree_nodes.append(node)
34.
35.       if highest_connection[node] != None:
36.         self.tree_edges.append((node, highest_connection[node]))
37.
38.       for other in self.q:
39.         if self.edges_cost[(other, node)] > cost_connection[other]:
40.           cost_connection[other] = self.edges_cost[(other, node)]
41.           highest_connection[other] = node
42.
43.     ################################
44.     # Transforming into directed tree
45.     edges_temp = self.tree_edges.copy()
46.     self.tree_root = self.tree_nodes[0]
47.
48.     for node in self.tree_nodes:
49.       edges_from_node = [edge for edge in edges_temp if node in edge]
50.       children = list(set([elem for edge in edges_from_node for elem in edge if elem != node]))
51.
52.       for edge in edges_from_node:
53.         edges_temp.remove(edge)
54.
55.       self.children_of[node] = children
56.
57.   ################################
58.   # returns the parent of a node
59.   def parent_of(self, child):
60.     for node in self.tree_nodes:
61.       if child in self.children_of[node]:
62.         return node
63.
64.     return None
65.
66.   def attr_parent_of(self, attr, obj):
67.     attr_field_name = attr.field_name
68.     parent_field_name = self.parent_of(attr_field_name)
69.     parent = None
70.
71.     for other in obj:
72.       if other.field_name == parent_field_name:
73.         parent = other
74.
75.     return parent
76.