迭代計算所有節點的子樹大小？

Question

我正在嘗試創建一個迭代版本：

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從堆棧中彈出，但是在那之后，我基本上在子循環之后“丟失”了ID，並且以后無法分配其大小。

我是否需要第二個堆棧？

Answer 1

未經測試-考慮此偽代碼的概念是要處理一堆節點，並且在每個節點上有相應的堆棧其直接子節點尚未處理。 這意味着主堆棧上的每個項目都是一個元組-元組中的第一個項目是節點，第二個項目是未處理的子節點列表。

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

編輯（尤其是測試后問題作者的編輯）當然是受歡迎的。