Build Railway Stations
Contest: Kotlin Heroes: Episode 9 (Unrated, T-Shirts + Prizes!)
Difficulty rating: 2000
Tags: *special, greedy, trees

Monocarp is playing a computer game where he's controlling an empire. An empire consists of $$$n$$$ cities, connected by $$$n - 1$$$ roads. The cities are numbered from $$$1$$$ to $$$n$$$. It's possible to reach every city from every other one using the roads.

Traversing every road takes $$$2$$$ hours. However, that can be improved. Monocarp can build railway stations in no more than $$$k$$$ cities. After they are built, all existing roads that connect two cities with railway stations get converted to railroads and become $$$1$$$ hour to traverse.

Let $$$f(x, y)$$$ be the total time it takes to traverse the roads on the shortest path between cities $$$x$$$ and $$$y$$$.

Monocarp wants to build at most $$$k$$$ railway stations in such a way that the following value is minimized: $$$\sum \limits_{v=1}^{n} \sum \limits_{u=1}^{v-1} f(v, u)$$$ (the total time it takes to travel from every city to every other one). What the smallest value he can achieve?

Input format:
The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of testcases.

The first line of each testcase contains two integers $$$n$$$ and $$$k$$$ ($$$2 \le k \le n \le 2 \cdot 10^5$$$) — the number of cities and the maximum number of railway stations Monocarp can build.

Each of the following $$$n-1$$$ lines contains two integers $$$v$$$ and $$$u$$$ ($$$1 \le v, u \le n$$$; $$$v \neq u$$$) — a road that connects cities $$$v$$$ and $$$u$$$.

It's possible to reach every city from every other one using the roads. The sum of $$$n$$$ over all testcases doesn't exceed $$$2 \cdot 10^5$$$.

Output format:
For each testcase, print a single integer — the smallest total time it takes to travel from every city to every other one that Monocarp can achieve after building at most $$$k$$$ railway stations.

Constraints:
Time limit: 3.0 s
Memory limit: 256.0 MB

Example input:
3
5 2
1 2
2 3
3 4
4 5
4 4
1 2
1 3
1 4
5 3
1 2
1 3
2 4
2 5

Example output:
34
9
26

Editorial:
Consider an easier problem. Instead of upgrading cities, let's upgrade the roads directly.
Since every path between two vertices on a tree is unique, upgrading a single road will decrease the answer by the number of paths that pass through it. These upgrades are completely independent of each other, so the solution would be to upgrade $$$k$$$ roads that have the most paths passing through them. This is easy to calculate. Root the tree by any vertex. For each vertex, calculate the number of vertices in its subtree. Then the value for an edge $$$(v, u)$$$, where $$$u$$$ is the lower vertex, is $$$\textit{size}_u \cdot (n - \textit{size}_u)$$$.
Turns out, this solution is basically the solution to the original problem. The edges with the most paths passing through them actually compose a connected subgraph of the tree. So it's always possible to choose $$$k-1$$$ best edges overall by picking some $$$k$$$ vertices.
How to prove it? Root the tree by its centroid. This will guarantee that the size of every subtree (except the root) is less than or equal to $$$\frac{n}{2}$$$. Consider some edge $$$(v, u)$$$. The idea is that for any child of $$$u$$$, an edge $$$(u, w)$$$ has lower value. Their numbers of paths are: $$$x \cdot (n - x)$$$ and $$$y \cdot (n - y)$$$ for some $$$x > y$$$. On the other hand, $$$\frac{n}{2} \ge x$$$. So it's common knowledge that it holds. Since going down the tree only decreases the number of paths, the edges with the most paths will always form a connected subgraph that includes the root — the centroid.
Overall complexity: $$$O(n \log n)$$$ or $$$O(n)$$$ (by coding nth_element function yourself) per testcase.