model-optimizer/mo/utils/dsu.py

   1 """
   2  Copyright (c) 2018-2019 Intel Corporation
   3
   4  Licensed under the Apache License, Version 2.0 (the "License");
   5  you may not use this file except in compliance with the License.
   6  You may obtain a copy of the License at
   7
   8       http://www.apache.org/licenses/LICENSE-2.0
   9
  10  Unless required by applicable law or agreed to in writing, software
  11  distributed under the License is distributed on an "AS IS" BASIS,
  12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  See the License for the specific language governing permissions and
  14  limitations under the License.
  15 """
  16
  17
  18 class DSUElem:
  19     """
  20     An object that represents one DSU element.
  21     """
  22     name = ''
  23     parent = ''
  24     rank = 1
  25
  26     def __init__(self, name):
  27         self.name = name
  28         self.parent = name
  29         self.rank = 1
  30
  31
  32 class DSU:
  33     """
  34     Naive implementation of the "disjoint set union" data structure.
  35     """
  36     map = dict()
  37
  38     def __init__(self, elems: list):
  39         self.map = {elem.name: elem for elem in elems}
  40         pass
  41
  42     def find_elem(self, name: str):
  43         return self.map[name]
  44
  45     def find_parent(self, elem: DSUElem):
  46         if elem.parent == elem.name:
  47             return elem
  48         parent_elem = self.find_parent(self.find_elem(elem.parent))
  49         elem.parent = parent_elem.name
  50         return parent_elem
  51
  52     def union(self, elem1: DSUElem, elem2: DSUElem):
  53         elem1 = self.find_parent(elem1)
  54         elem2 = self.find_parent(elem2)
  55         if elem1.name == elem2.name:  # already in the same set
  56             return
  57
  58         if elem1.rank < elem2.rank:
  59             elem1.parent = elem2.name
  60         elif elem1.rank > elem2.rank:
  61             elem2.parent = elem1.name
  62         else:
  63             elem1.parent = elem2.name
  64             elem2.rank = elem2.rank + 1