37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
from obitools.graph import UndirectedGraph,Node
|
|
from obitools.graph.algorithms.component import componentCount
|
|
|
|
|
|
class Forest(UndirectedGraph):
|
|
|
|
|
|
def getNode(self,node=None,index=None):
|
|
if index is None:
|
|
index = self._index.getIndex(node, True)
|
|
return TreeNode(index,self)
|
|
|
|
def addEdge(self,node1=None,node2=None,index1=None,index2=None,**data):
|
|
index1=self.addNode(node1, index1)
|
|
index2=self.addNode(node2, index2)
|
|
|
|
cc = set(n.index for n in self.getNode(index=index2).componentIterator())
|
|
|
|
assert index1 in self._node[index2] or index1 not in cc, \
|
|
"No more than one path is alloed between two nodes in a tree"
|
|
|
|
UndirectedGraph.addEdge(self, index1=index1, index2=index2,**data)
|
|
|
|
return (index1,index2)
|
|
|
|
def isASingleTree(self):
|
|
return componentCount(self)==1
|
|
|
|
class TreeNode(Node):
|
|
|
|
def componentIterator(self):
|
|
for c in self:
|
|
yield c
|
|
for cc in c:
|
|
yield cc
|
|
|
|
|