Advertisement
Fabio_LaF

Classe MaxTree

Aug 17th, 2022
554
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 2.40 KB | None | 0 0
  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.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement