#####   ######  #####    ###    #   #  ###  #   #  ######
##  ##  ##      ##  ##  ## ##   #   #   #   #   #  ##
#####   ####    #####   #   #   #   #   #   #   #  ####
##  #   ##      ##      ## ##    # #    #    # #   ##
##   #  ######  ##       ###      #    ###    #    ######

$ curl repovive.com/roadmaps/graph-theory/tree-diameter-center/tree-distances-ii-implementation

░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░███████████████████████████████████████████████████████████████████████████████████████████████

#####   ######  #####    ###    #   #  ###  #   #  ######
##  ##  ##      ##  ##  ## ##   #   #   #   #   #  ##
#####   ####    #####   #   #   #   #   #   #   #  ####
##  #   ##      ##      ## ##    # #    #    # #   ##
##   #  ######  ##       ###      #    ###    #    ######

$ curl repovive.com/roadmaps/graph-theory/tree-diameter-center/tree-distances-ii-implementation

Repovive

CSES 1133 Tree Distances II - Implementation - Tree Diameter & Center | Graph Theory | Repovive

Repovive.

Graph Theory37 sections · 1633 units81h total

Open in Course

CSES 1133 Tree Distances II - Implementation

Python code

Here is the solution:

subtreeSum = array of size n + 1
subtreeSize = array of size n + 1
ans = array of size n + 1

function dfsDown(v, parent):
    subtreeSize[v] = 1
    subtreeSum[v] = 0
    for u in adj[v]:
        if u != parent:
            dfsDown(u, v)
            subtreeSize[v] = subtreeSize[v] + subtreeSize[u]
            subtreeSum[v] = subtreeSum[v] + subtreeSum[u] + subtreeSize[u]

function dfsUp(v, parent, total):
    ans[v] = total
    for u in adj[v]:
        if u != parent:
            dfsUp(u, v, total - subtreeSize[u] + (n - subtreeSize[u]))

dfsDown(1, 0)
dfsUp(1, 0, subtreeSum[1])

This runs in $O(n)$ time and uses $O(n)$ space.