SubMST

We need the sum of MST weights over all subsets in the complete graph where edge weights are tree distances.

First, murder the tempting false shortcut: the answer for a subset is not always the size of the Steiner tree in the original tree. On a star with three selected leaves, the Steiner tree has weight $3$, but the metric MST has two edges of weight $2$, total $4$. So yeah, that idea gets packed up and sent home.

The Kruskal Decomposition

Take one fixed subset $S$. Consider its MST in the complete distance graph.

For a tree metric, every chosen MST edge $(a,b)$ corresponds to the path from $a$ to $b$ in the original tree. That path has a middle. The middle is either:

• a vertex, if $\operatorname{dist}(a,b)$ is even;
• an original tree edge, if $\operatorname{dist}(a,b)$ is odd.

So the MST weight can be decomposed by these centers.

It is enough to count contributions from every center independently, then add them all.

Vertex Centers

Fix a vertex $r$.

Delete $r$ from the tree. The remaining connected components are the branches around $r$. Also treat $r$ itself as a special branch containing one vertex at distance $0$.

For a subset $S$, each branch is either empty or has a closest selected vertex to $r$. Let that closest distance be $d_B$.

If two selected vertices lie in different branches around $r$, their path goes through $r$, and the cheapest connection between those two branches has weight

$$d_A+d_B.$$

So for this center $r$, we reduce the subset to a bunch of nonempty branches with values $d_1,d_2,\ldots,d_k$.

The MST of the complete graph on these branches with edge weights $d_i+d_j$ is easy. If

$$d_1 \le d_2 \le \cdots \le d_k,$$

then the optimal tree connects every other branch to the smallest one, giving

$$\sum_{i=2}^{k}(d_1+d_i).$$

Proof is blunt: branch $i>1$ must connect to something, and every edge incident to it costs at least $d_1+d_i$. The star through the minimum branch achieves exactly that lower bound.

Edge Centers

Now handle centers that are original tree edges.

For an edge $(u,v)$, removing it splits the tree into two sides. If the closest selected vertex on the $u$ side is distance $x$ from $u$, and the closest selected vertex on the $v$ side is distance $y$ from $v$, then the cheapest connection whose middle is this edge has weight

$$x+1+y.$$

This is the same idea as vertex centers, except there are only two sides and the middle edge contributes $1$.

So every MST edge is counted exactly once: either by its middle vertex or by its middle tree edge.

Counting Choices

For a fixed center and one side/branch, suppose it contains $cnt[d]$ vertices at distance $d$ from the center endpoint.

The number of ways for the nearest selected vertex in that branch to be exactly distance $d$ is

$$ways[d]=(2^{cnt[d]}-1)\cdot 2^{\sum_{e>d}cnt[e]}.$$

We choose at least one vertex at distance $d$, choose no vertices closer, and choose anything farther. Empty branch has one way.

DP for a Vertex Center

For a vertex center $r$, process its branches one by one.

Maintain:

• $dp[m]$: number of processed selections where at least one processed branch is nonempty and the smallest nearest distance is $m$;
• $sum[m]$: total contribution already accumulated for those selections.

When adding a new nonempty branch with nearest distance $x$:

• if this is the first nonempty branch, contribution is $0$;
• otherwise, it connects to the current best branch and adds $m+x$;
• the new minimum becomes $\min(m,x)$.

This incremental process produces exactly

$$\sum_{i=2}^{k}(d_1+d_i)$$

for every chosen set of nonempty branches, regardless of processing order. Empty branch just leaves the state unchanged.

DP for an Edge Center

For an edge center $(u,v)$, there are only two sides. Count all choices where the nearest selected vertex on the $u$ side has distance $x$, and on the $v$ side has distance $y$. Their contribution is

$$x+1+y.$$

Multiply the two side counts and sum.

Implementation Details

For every root/center, run BFS/DFS to get distances. Since $n\le 5000$ and $\sum n^2\le 25\cdot 10^6$, doing $O(n)$ work per center is fine.

For vertex centers, BFS from $r$ also identifies the first neighbor of $r$ on the path to every vertex, which tells us the branch.

For edge centers, root the tree once per directed edge side by a DFS that avoids crossing the center edge, collecting distance counts.

This gives an $O(n^2)$ solution overall if implemented with sane constants.

Correctness Proof

First, every edge in the metric MST corresponds to a path in the original tree. That path has a unique middle: either a vertex or an original tree edge. Therefore, assigning every MST edge to its middle partitions the MST weight into independent center contributions.

For a fixed vertex center $r$, selected vertices from the same branch around $r$ have paths whose middle is not $r$, so they do not contribute to center $r$. Selected vertices from different branches have paths through $r$, and the cheapest possible edge between two such branches uses the closest selected vertex in each branch, with cost $d_A+d_B$.

Thus the center-$r$ contribution only depends on the nearest selected distance of every nonempty branch. The induced branch graph has edge weights $d_i+d_j$. Its MST is the star through the smallest $d_i$, with weight $\sum_{i=2}^k(d_1+d_i)$, because each non-minimum branch needs an incident edge costing at least $d_1+d_i$, and the star achieves all those bounds.

The branch DP computes exactly this value for every subset: when a new branch with nearest distance $x$ is added to an existing collection with minimum $m$, the new branch contributes the cheapest necessary connection $m+x$, and the new minimum becomes $\min(m,x)$. Empty branches contribute nothing. Since ways[d] counts exactly the selections with nearest distance $d$, every subset is counted once.

For a fixed edge center $(u,v)`, removing the edge gives two sides. A path has this edge as its middle exactly when it connects selected vertices from opposite sides with closest distances$x$and$y$from the endpoints, contributing$x+1+y$. The formula sums exactly those choices.

Since every MST edge has exactly one middle, and each center contribution is counted exactly once, summing all vertex-center and edge-center contributions over all subsets gives the required total.

Complexity

The algorithm does $O(n)$ work for each of $O(n)$ centers, so the total time is $O(n^2)$ per test case. Memory is $O(n)$ besides the adjacency list.