迭代计算所有节点的子树大小?

时间:2015-08-20 18:22:08

标签: python algorithm graph tree depth-first-search

我正在尝试创建一个迭代版本:

def computeSize(id):
   subtreeSize[id] = 1
   for child in children[id]:
      computeSize(child)
      subtreeSize[id]+=subtreeSize[child]

"迭代"意思是没有递归,因为在Python中,如果你的图形很大并且在任何地方都有冗长的线性链,它会给出一个堆栈递归错误。

尝试使用堆栈代替(使用DFS算法对其进行建模)但我对细节有困难:

def computeSubtreeSizes(self): #self.sizes[nodeID] has size of subtree
    stack = [self.rootID] #e.g. rootID = 1
    visited = set()

    while stack:
        nodeID = stack.pop()
        if nodeID not in visited:
            visited.add(nodeID)
            for nextNodeID in self.nodes[nodeID]:
                stack.append(nextNodeID)

例如,一旦我开始,我会明显地将根ID从堆栈中弹出,但在那之后,我基本上已经#34;丢失"子循环后的ID,以后无法分配其大小。

我是否需要第二个堆栈?

1 个答案:

答案 0 :(得分:2)

未经测试 - 将此伪代码视为处理一堆节点的概念,并在每个节点上考虑尚未处理的直接子节点的相应堆栈。这意味着主堆栈上的每个项目都是一个元组 - 元组中的第一项是节点,第二项是未处理的子节点列表。

def computeSubtreeSizes(self):
    stack = [(self.rootID, [])] #e.g. rootID = 1
    visited = self.sizes = {}

    while stack:
        nodeID, subnodes = stack[-1]
        size = visited.get(nodeID)
        if size is None:
            # Haven't seen it before.  Set total to 1,
            # and set up the list of subnodes.
            visited[nodeID] = size = 1
            subnodes[:] = self.nodes[nodeID]
        if subnodes:
            # Process all the subnodes one by one
            stack.append((subnodes.pop(), []))
        else:
            # When finished, update the parent
            stack.pop()
            if stack:
                visited[stack[-1][0]] += size

显而易见的潜在性能提升是不打扰将已访问过的节点添加到主堆栈。 这仅在重复的子树非常常见时才有用。这是更多的代码(可读性更低),但可能看起来像这样:

def computeSubtreeSizes(self):
    stack = [(self.rootID, [])] #e.g. rootID = 1
    visited = self.sizes = {}

    while stack:
        nodeID, subnodes = stack[-1]
        size = visited.get(nodeID)
        if size is None:
            # Haven't seen it before.  Add totals of
            # all previously visited subnodes, and
            # add the others to the list of nodes to
            # be visited.
            size = 1
            for sn in self.nodes[nodeID]:
                sn_size = visited.get(sn)
                if sn_size is None:
                    subnodes.append(sn)
                else:
                    size += sn_size
            visited[nodeID] = size

        if subnodes:
            # Process all the subnodes one by one
            stack.append((subnodes.pop(), []))
        else:
            # When finished, update the parent
            stack.pop()
            if stack:
                visited[stack[-1][0]] += size

编辑(特别是测试后的问题作者)当然是受欢迎的。