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

Iteratively compute subtree sizes for all nodes?

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

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

"Iterative" 意味着没有递归,因为在 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 从堆栈中弹出,但在那之后,我基本上 "lost" 子循环之后的 ID 并且以后无法分配它的大小.

我需要第二个堆栈吗?

未测试 -- 考虑这个伪代码的概念是有一堆正在处理的节点,并且在每个节点上,一个对应的尚未处理的直接子节点堆栈。这意味着主堆栈上的每一项都是一个元组——元组中的第一项是节点,第二项是未处理的子节点列表。

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

当然欢迎编辑(尤其是测试后问题作者的编辑)。